from ca import CASim
import numpy as np
import matplotlib.pyplot as plt
import json
import os.path
import scipy.interpolate as sint
import scipy.optimize as sop
import pickle


def run_sim(N=50, T=50, dens=0.0, plot=False, fig_name="", ret_flow=False):
    """Runs the car simulation with the given parameters"""
    ca = CASim(w=N, h=T, c_dens=dens)
    ca.reset()
    for i in range(T):
        ca.step()
    if plot:
        ca.draw(f"Run N={N} for {T} time steps with car density {dens}", fig_name)
    if ret_flow:
        return ca.flow / T

def int_crit(flow, density):
    """Interpolates the critical density"""
    f = sint.interp1d(density, flow, kind='cubic')
    f_max = sop.minimize_scalar(lambda x: -f(x), method='bounded', bounds=(0,1))
    return f_max.x

def main():
    ## 2
    # Run N=50 for 50 time steps with c_dens 0.4
    flow = run_sim(dens=0.4, plot=True, fig_name="fig2.1.png", ret_flow=True)

    # Run N=50 for 50 time steps with c_dens 0.9
    flow = run_sim(dens=0.9, plot=True, fig_name="fig2.2.png", ret_flow=True)


    ## 4
    R = 10
    d_val = np.linspace(0, 1, 21)
    mean_flow = []
    if not os.path.isfile("data4.json"):
        for d in d_val:
            print(d)
            f_sum = 0
            for i in range(int(R)):
                f_sum += run_sim(N=50, T=1000, dens=d, plot=False, ret_flow=True)
            mean_flow.append(f_sum / R)

        data4 = dict(dens=list(d_val), flow=mean_flow)
        f = open("data4.json", 'w')
        json.dump(data4, f)
    else:
        f = open("data4.json", 'r')
        data4 = json.load(f)
        d_val = data4['dens']
        mean_flow = data4['flow']

    plt.clf()
    plt.scatter(d_val, mean_flow)
    plt.title(f"Mean flow for {R} runs with {len(d_val)} density values for N = 50 and T = 1000")
    plt.xlabel("Density")
    plt.ylabel("Average flow per time step")
    plt.savefig("fig4.png")
    plt.clf()


    ## 5
    R = 3
    d_val = np.linspace(0, 1, 21)
    mean_flow = []
    if not os.path.isfile("data5.json"):
        for d in d_val:
            print(d)
            f_sum = 0
            for i in range(int(R)):
                f_sum += run_sim(N=50, T=5, dens=d, plot=False, ret_flow=True)
            mean_flow.append(f_sum / R)

        data5 = dict(dens=list(d_val), flow=mean_flow)
        f = open("data5.json", 'w')
        json.dump(data5, f)
    else:
        f = open("data5.json", 'r')
        data5 = json.load(f)
        d_val = data5['dens']
        mean_flow = data5['flow']

    plt.clf()
    plt.scatter(d_val, mean_flow)
    plt.title(f"Mean flow for {R} runs with {len(d_val)} density values for N = 50 and T = 5")
    plt.xlabel("Density")
    plt.ylabel("Average flow per time step")
    plt.savefig("fig5.png")
    plt.clf()


    ## 6
    crit4 = int_crit(data4['flow'], data4['dens'])
    d_val = np.linspace(0, 1, 21)
    T_val = np.arange(1, 52, 5)
    R = 10
    if not os.path.isfile("data6.pickle"):
        n_T_cor= []
        for T in T_val:
            #TODO
            print(T)
            cor = 0
            for i in range(20):
                mean_flow = []
                for d in d_val:
                    f_sum = 0
                    for i in range(R):
                        f_sum += run_sim(N=50, T=T, dens=d, plot=False, ret_flow=True)
                    mean_flow.append(f_sum / R)
                crit = int_crit(mean_flow, d_val)
                print(f"T: {T}, crit: {crit}")
                if abs(crit - 0.5) <= 0.05:
                    cor += 1
            n_T_cor.append(cor)
        data6 = dict(T=list(T_val), cor=n_T_cor)
        f = open("data6.pickle", 'wb')
        pickle.dump(data6, f)
    else:
        data6 = pickle.load(open("data6.pickle", 'rb'))
        T_val = data6['T']
        n_T_cor = data6['cor']

    print(n_T_cor)

    cor_perc = [x/20 * 100 for x in n_T_cor]

    plt.clf()
    plt.scatter(T_val, cor_perc)
    plt.ylim([0, 110])
    plt.title(f"Percentage of runs with correct critical density per time window T")
    plt.xlabel("T")
    plt.ylabel("Percentage with correct critical density")
    plt.savefig("fig6.png")

if __name__ == "__main__":
    main()
