import unittest
import numpy as np
from analysis import Analysis
import matplotlib.pyplot as plt

# Parameter to set how far the results can deviate from the true value
# (We use much larger errors here for the initial data than for the
# simualtion, so we allow for a bigger range. We are randomize error
# as oppodes to the actial simualtion.)
sigma = 10

class testAnalysis(unittest.TestCase):
    """
    A simple test to check the correct results for individual fits
    and for the combined fit.
    The individual fits can only be tested upto absolute values as for
    the correct sign determination the information of all fits is
    required.
    """

    def testxFit(self):
        A = Analysis()
        #initial parameters
        param = [35] #Vx

        t = np.arange(0.3, 0.5, 0.025)
        x = A.xEOM(param, t)

        Data = []
        for i in range(0,len(t)):
            err =0.01* np.random.random(size=(2))
            Data.append([[x[i]*(1+err[0]),x[i]*(1-err[1])],
                         [0,0],
                         [0,0],
                         [t[i], t[i]]])
        Data = np.array(Data)

        A.fill(Data)
        OUT = A.xFit()
        self.assertTrue(np.abs(param[0]) < np.abs(OUT.beta[0]) + sigma*OUT.sd_beta[0])
        self.assertTrue(np.abs(param[0]) > np.abs(OUT.beta[0]) - sigma*OUT.sd_beta[0])


    def testyFit(self):
        A = Analysis()
        param = [2, 25, 250]

        t = np.arange(0.3, 0.5, 0.025)
        y = A.yEOM(param, t)
        
        Data = []
        for i in range(0,len(t)):
            err =0.01* np.random.random(size=(2))
            Data.append([[0,0],
                         [y[i]*(1+err[0]),y[i]*(1-err[1])],
                         [0,0],
                         [t[i], t[i]]])
        Data = np.array(Data)

        A.fill(Data)
        OUT = A.yFit()
        self.assertTrue(np.abs(param[0]) < np.abs(OUT.beta[0]) + sigma*OUT.sd_beta[0])
        self.assertTrue(np.abs(param[0]) > np.abs(OUT.beta[0]) - sigma*OUT.sd_beta[0])


    def testzFit(self):
        A = Analysis()
        param = [2, 25, 250]

        t = np.arange(0.3, 0.5, 0.025)
        y = A.zEOM(param, t)

        Data = []
        for i in range(0,len(t)):
            err =0.01* np.random.random(size=(2))
            Data.append([[0,0],
                         [0,0],
                         [y[i]*(1+err[0]),y[i]*(1-err[1])],
                         [t[i], t[i]]])
        Data = np.array(Data)

        A.fill(Data)
        OUT = A.zFit()
        self.assertTrue(np.abs(param[0]) < np.abs(OUT.beta[0]) + sigma*OUT.sd_beta[0])
        self.assertTrue(np.abs(param[0]) > np.abs(OUT.beta[0]) - sigma*OUT.sd_beta[0])


    def testFullFit(self):
        ANA = Analysis()
        w = -2
        vx = 35
        vy = -25
        vz = 250

        t = np.arange(0.3, 0.5, 0.025)
        x = ANA.xEOM([vx], t)
        y = ANA.yEOM([w,vy,vz], t)
        z = ANA.zEOM([w,vy,vz], t)
        
        Data = []
        for i in range(0,len(t)):
            err =0.01* np.random.random(size=(6))
            Data.append([[x[i]*(1+err[0]),x[i]*(1-err[1])],
                         [y[i]*(1+err[2]),y[i]*(1-err[3])],
                         [z[i]*(1+err[4]),z[i]*(1-err[4])],
                         [t[i], t[i]]])
        Data = np.array(Data)

        ANA.fill(Data)
        ANA.Fit()
        ANA.output()

        self.assertTrue(w  < ANA.results['w'][0]  + sigma*ANA.results['w'][1])
        self.assertTrue(w  > ANA.results['w'][0]  - sigma*ANA.results['w'][1])
        self.assertTrue(vx < ANA.results['vx'][0] + sigma*ANA.results['vx'][1])
        self.assertTrue(vx > ANA.results['vx'][0] - sigma*ANA.results['vx'][1])
        self.assertTrue(vy < ANA.results['vy'][0] + sigma*ANA.results['vy'][1])
        self.assertTrue(vy > ANA.results['vy'][0] - sigma*ANA.results['vy'][1])
        self.assertTrue(vz < ANA.results['vz'][0] + sigma*ANA.results['vz'][1])
        self.assertTrue(vz > ANA.results['vz'][0] - sigma*ANA.results['vz'][1])


if __name__ == "__main__":
    unittest.main()