from numpy import sqrt, arcsin, arccos, cos, sin

from scipy.optimize import minimize
import numpy

from Params import Params



class OptTimeCalculator():
    
    def __init__(self, params : Params):
        self.Ts = {}
        self.params = params
        self.Tmax = 1
        
    def tGrow(self, F):
        return arcsin(F/(self.params.Ftg)) * sqrt(self.params.m/self.params.k)

    def T(self, s, s2 = 0, debug = False):
        if hasattr(s, "__len__"):
            s = s[0]
        v0q = min(sqrt(2 * self.params.a * s), self.params.v1limit)
        v0 = min(v0q, sqrt(1/(self.params.k*self.params.m))* self.params.Ftg)
        t1 = v0 / self.params.a
        
        #test second
        t21 = sqrt(s2/self.params.a2)
        t21 = min(self.params.v2limit/self.params.a2, t21)
        t22 = max(0, (s2 - (self.params.a2 * t21 * t21)) / self.params.v2limit)
        T2 = t22 + 2 * t21
        
        t2t = max(0, (s - (self.params.a * t1 * t1 /2)) / v0q)
        #if debug:
            #print("!", t1, t2t, "|", t21, t22, "|", t1+t2t, 2*t21+t22, t2t + t1 < T2)
        if t2t + t1 < T2:
            if debug:
                print("bad time")
            return 2 * self.Tmax
        #print("t2", t2t) 
        s3 = s+((self.params.Fd - self.params.fl)/2 + self.params.fl)/self.params.k
        t3 = sqrt(s3 / self.params.a)
        v3max = min(t3 * self.params.a, self.params.v1limit)
        t3 = v3max / self.params.a
        t3q = max(0, (s3 - t3*t3*self.params.a)/v3max)
        v0 = v0 * self.params.k
        l = sqrt(self.params.awork**2 * self.params.m**2 + self.params.m/self.params.k * v0**2)
        #print("v0", v0, v0/k)
        Fmeet = (self.params.Ftg**2 - self.params.m/self.params.k*v0**2)/(2 * self.params.awork * self.params.m)
        print(Fmeet)
        #if debug:
            #print(Fmeet)
        Fq = self.params.Fq
        if Fmeet > Fq:
            if debug:
                print("cant reach", Fmeet)
            return self.Tmax
        #print(Ftg - Fmeet)
        if Fmeet - self.params.awork*self.params.m > l:
            if debug:
                print("wut")
            return 3*self.Tmax
        beta = self.params.awork * self.params.m
        t2 = sqrt(self.params.m/self.params.k) * (arcsin((Fmeet - beta)/(l)) + arccos(sqrt(self.params.m/self.params.k) * v0/l))
        if debug:
            qtq = sqrt(self.params.k/self.params.m)
            Fqq = -self.params.awork*self.params.m*cos(qtq*t2) + self.params.awork*self.params.m + v0 / qtq * sin(qtq*t2)
            print("angles", (Fmeet / (2*beta) -beta)/(l), (sqrt(self.params.m/self.params.k) * v0/l), Fqq, Fmeet)
        tend = self.tGrow(Fq) - self.tGrow(Fmeet)
        vp = 1/sqrt(self.params.k * self.params.m) * sqrt(self.params.Ftg**2 - self.params.Fq**2)
        ap = Fq / self.params.m
        tprop = 2*vp / ap
        #print(s, t1, t2, Fmeet, tend)
        T = t1 + t2t + t2 + 2*t3 + t3q + tprop + tend
        
        self.Ts = {"t1":t1, "t2":t2t, "tm":t2, "tend":tend, "t4":tprop, "to1":t3, "to2": t3q, "T":T,
                   "tclose":t1+t2t, "tgrow":t2+tend+tprop, "topen": 2*t3+t3q}
        if debug:
            print("phases", t1, t2t, "|", t2, tend, tprop, "|", 2*t3+t3q, "|", T)
            print("tclose: ",t1+t2t, "tgrow:", t2+tend+tprop, "topen:", 2*t3+t3q, "all:", T)
        return T
    
    def calcSecond(self, Tstart, s2):
        T1 = self.Ts["tclose"]
        t1 = T1 / 2 - sqrt(T1**2 - 4 * s2 / self.params.a2) / 2
        t2 = sqrt(T1 * T1 - 4 * s2 / self.params.a2)
        
        T2 = self.Ts["topen"] - Tstart
        to1 = T2 / 2 - sqrt(T2**2 - 4 * s2 / self.params.a2) / 2
        to2 = sqrt(T2 * T2 - 4 * s2 / self.params.a2)
        
        Ts = {"t1":t1, "t2":t2, "to1":to1, "to2": to2}
        self.T2 = Ts
        return self.T2

    def calcOpt(self):
        s0 = [0.04]
        bnds = [(0.00001, None)]
        res = minimize(self.T, s0, method='Nelder-Mead', tol=1e-7, bounds=bnds)
        nonbound = res.x
        bound = res.x
        if nonbound < self.params.smin1:
            #print("IS SMIN")
            bound = max(nonbound, self.params.smin1)
        #T(q, debug = True)
        return nonbound, bound

    def showPlot(self):

        N = 50000
        ss = numpy.linspace(0.0001, 0.030, N)
        Ts = numpy.array([self.T(si) for si in ss])

        import matplotlib.pyplot as plt
        from matplotlib import use

        q, qq  = self.calcOpt()
        q = 1000*q
        qq = 1000 * qq
        use("GTK4Agg", force = True)
        plt.plot(ss*1000, Ts*1000, "red", label = "T(S)")
        plt.plot(numpy.ones(N) * qq, Ts*1000, label = "factOpt")
        plt.plot(numpy.ones(N) * q, Ts*1000, "black", label = "trueOpt")
        plt.title("T(s)")
        plt.xlabel("s, мм")
        plt.ylabel("T, мс")
        plt.legend()
        plt.savefig("model.png")
        plt.show()
        
    def getTimes(self, s):
        self.T(s, debug=False)
        return self.Ts
   

if __name__ == "__main__":
    opt = OptTimeCalculator() 
    nonbound, bound = opt.calcOpt()
    
    print("nonbound", nonbound*1000, "мм")
    print(opt.getTimes(nonbound))
    
    print("bound", bound*1000, "мм")
    print(opt.getTimes(bound))