import sys
import plot
import subprocess
import math as mth
import cmath as cmt
import numpy as npy
import scipy.linalg as sla

state_num = 1
dt = 0.02
accu = 100
iterations = 5000


M0 = [[1,0,0,0],[0,1,0,0],[0,0,-1,0],[0,0,0,-1]]
M1 = [[0,0,1,0],[0,0,0,1],[1,0,0,0],[0,1,0,0]]
M2 = [[0,0,-1j,0],[0,0,0,-1j],[1j,0,0,0],[0,1j,0,0]]
M3 = [[0,0,0,1],[0,0,-1,0],[0,-1,0,0],[1,0,0,0]]


def init_state(i):
    alpha = 2
    beta  = 0
    gamma = 1
    delta = 1+1j
    if state_num > 1:
        alpha = i / state_num
        beta  = -(state_num - i) / state_num
        gamma = alpha + beta
        delta = delta * i
    H = npy.array([[1,0,0,1j],[0,2,0,0],[0,0,3,0],[-1j,0,0,4]])
    H = sla.expm(-1j * (dt / accu) * H)
    norm  = npy.linalg.norm([alpha, beta, gamma, delta])
    state = npy.array([alpha / norm, beta / norm, gamma / norm, delta / norm])
    return (state, H)


def time_evolution(state, dt=dt):
    return (npy.dot(*state), state[1])

def fibration(state):
    x0=npy.dot(npy.conj(state),npy.dot(M0,state))
    x1=npy.dot(npy.conj(state),npy.dot(M1,state))
    x2=npy.dot(npy.conj(state),npy.dot(M2,state))
    V=npy.dot(state,npy.dot(M3,state))
    x3=V.real
    x4=V.imag

    return ([x0.real,x1.real,x2.real,x3,x4])


states = []
for i in range(state_num):
    (co, H) = init_state(i)
    norm  = npy.linalg.norm(co)
    state = npy.array([co[0] / norm, co[1] / norm, co[2] / norm, co[3] / norm])
    states.append((state, H))

f = open("data", "w")
for i in range(iterations):
    for j in range(state_num):
        hopf_state = fibration(states[j][0])
        colour = i / iterations
        if state_num > 1:
            colour = j / state_num
        f.write(f"{hopf_state[0]}; {hopf_state[1]}; {hopf_state[2]}; {hopf_state[3]}; {hopf_state[4]}; {colour}\n")
        for l in range(accu):
            states[j] = time_evolution(states[j], dt / accu)
f.close()

plot.plot(iterations, iterations, state_num, "anim3d.plt")
#plot.plot(iterations, iterations, state_num, "anim2d.plt")