import numpy as np

from josephson_circuit import *
from time_evolution import *

import matplotlib
import matplotlib.pyplot as plt
matplotlib.use("TkAgg")

"""
Time evolution example 4: Hysteresis

Ramp current up and down to see if the voltage is hysteretic.

Done with and without capacitance and inductance. 

"""

# define array
N = 20
sq_array_no_scr_no_cap = SquareArray(N, N)

sq_array_scr_no_cap = SquareArray(N, N)
sq_array_scr_no_cap.set_inductance_factors(2)

sq_array_no_scr_cap = SquareArray(N, N)
sq_array_no_scr_cap.set_capacitance_factors(3)

sq_array_scr_cap = SquareArray(N, N)
sq_array_scr_cap.set_inductance_factors(2)
sq_array_scr_cap.set_capacitance_factors(3)


# define problem
T = 0.05
dt = 0.05
Nt = 30000
t_gap = 200
ts = np.arange(0, 2 * Nt, t_gap)
Imin = 0.4
Imax = 1.6
Ih = sq_array_no_scr_no_cap.horizontal_junctions()
It = np.append(np.linspace(Imin, Imax, Nt), np.linspace(Imax, Imin, Nt))
Is = Ih[:, None, None] * It
prob_no_scr_no_cap = TimeEvolutionProblem(sq_array_no_scr_no_cap, time_step=dt, time_step_count=2*Nt, current_sources=Is, temperature=T, store_time_steps=ts)
prob_scr_no_cap = TimeEvolutionProblem(sq_array_scr_no_cap, time_step=dt, time_step_count=2*Nt, current_sources=Is, temperature=T, store_time_steps=ts)
prob_no_scr_cap = TimeEvolutionProblem(sq_array_no_scr_cap, time_step=dt, time_step_count=2*Nt, current_sources=Is, temperature=T, store_time_steps=ts)
prob_scr_cap = TimeEvolutionProblem(sq_array_scr_cap, time_step=dt, time_step_count=2*Nt, current_sources=Is, temperature=T, store_time_steps=ts)

# do time simulation
out_no_scr_no_cap = prob_no_scr_no_cap.compute()
out_scr_no_cap = prob_scr_no_cap.compute()
out_no_scr_cap = prob_no_scr_cap.compute()
out_scr_cap = prob_scr_cap.compute()

def get_V(th):
    return np.append([0], np.mean(np.diff(th[Ih, 0, :], axis=-1) / (dt * t_gap), axis=0))


# compute array voltage
V_no_scr_no_cap = get_V(out_no_scr_no_cap.get_theta())
V_scr_no_cap = get_V(out_scr_no_cap.get_theta())
V_no_scr_cap = get_V(out_no_scr_cap.get_theta())
V_scr_cap = get_V(out_scr_cap.get_theta())

# plot array voltage
plt.plot(It[ts], V_no_scr_no_cap, label='betaL=0, betaC=0')
plt.plot(It[ts], V_scr_no_cap, label='betaL=2, betaC=0')
plt.plot(It[ts], V_no_scr_cap, label='betaL=0, betaC=2')
plt.plot(It[ts], V_scr_cap, label='betaL=2, betaC=2')
plt.xlabel("current")
plt.ylabel("voltage")
plt.title("Hysteresis of IV curve with or without screening and capacitance")
plt.legend()
plt.show()
