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

class testAnalysis(unittest.TestCase):
    """
    A simple test to check the angle means ans error computation.
    """

    def testxFit(self):
        A = Analysis()
        param = [35]

        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(param[0] < OUT.beta[0] + 3*OUT.sd_beta[0])
        self.assertTrue(param[0] > OUT.beta[0] - 3*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(param[0] < OUT.beta[0] + 3*OUT.sd_beta[0])
        self.assertTrue(param[0] > OUT.beta[0] - 3*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(param[0] < OUT.beta[0] + 3*OUT.sd_beta[0])
        self.assertTrue(param[0] > OUT.beta[0] - 3*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()
        
        self.assertTrue(w  < ANA.results['w'][0]  + 3*ANA.results['w'][1])
        self.assertTrue(w  > ANA.results['w'][0]  - 3*ANA.results['w'][1])
        self.assertTrue(vx < ANA.results['vx'][0] + 3*ANA.results['vx'][1])
        self.assertTrue(vx > ANA.results['vx'][0] - 3*ANA.results['vx'][1])
        self.assertTrue(vy < ANA.results['vy'][0] + 3*ANA.results['vy'][1])
        self.assertTrue(vy > ANA.results['vy'][0] - 3*ANA.results['vy'][1])
        self.assertTrue(vz < ANA.results['vz'][0] + 3*ANA.results['vz'][1])
        self.assertTrue(vz > ANA.results['vz'][0] - 3*ANA.results['vz'][1])