# Copyright 2018 Raytheon BBN Technologies
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0

import unittest

import os
import numpy as np
from auspex.analysis import fits, qubit_fits, resonator_fits

import matplotlib.pyplot as plt


class FitAssertion(object):

    def assertFitInterval(self, p0, name, fit, tol=5):
        low = fit.fit_params[name] - tol*fit.fit_errors[name]
        high = fit.fit_params[name] + tol*fit.fit_errors[name]
        test =  (low < p0 < high)
        if not test:
            raise AssertionError(f"Fit parameter {name}: {p0} is outside of interval ({low}, {high}).")

class TestResonatorFit(unittest.TestCase, FitAssertion):

    def test_CircleFit(self):

        #[tau, a, alpha, fr, phi0, Ql, Qc, Qi]

        Qi = 6.23e5
        Qc = 2e5
        Ql = 1/(1/Qi + np.real(1/Qc))
        f0 = 6.86
        kappa = f0/Ql

        p0 = [(1/f0)*0.9734, 0.8, np.pi*0.09, f0, np.pi*0.123, Ql, Qc]

        x = np.linspace(f0 - 8*kappa, f0+7*kappa, 1601)
        y = resonator_fits.ResonatorCircleFit._model(x, *p0)

        noise = 1.0 + np.random.randn(y.size) * np.median(y)/20
        y *= noise

        fit = resonator_fits.ResonatorCircleFit(x, y, make_plots=False)

        #print(fit.fit_params)
        #print(fit.fit_errors)

        try:
            self.assertFitInterval(f0, "fr", fit, tol=10)
            self.assertFitInterval(Qi, "Qi", fit, tol=2)
            self.assertFitInterval(Ql, "Ql", fit, tol=2)
        except AssertionError as e:
            print("Resonator fit tests failed. Perhaps re-run?")
            print(str(e))
        except:
            pass


class TestFitMethods(unittest.TestCase, FitAssertion):

    def test_LorentzFit(self):

        p0 = [3, 0.25, 0.4, 1.0]
        x = np.linspace(-1, 1, 201)
        y = fits.LorentzFit._model(x, *p0)
        noise = np.random.randn(y.size) * np.max(y)/10
        y += noise
        fit = fits.LorentzFit(x, y, make_plots=False)
        self.assertFitInterval(p0[0], "A", fit)
        self.assertFitInterval(p0[1], "b", fit)
        self.assertFitInterval(p0[2], "c", fit)

    def test_GaussianFit(self):
        p0 = [0.23, 3.1, 0.54, 0.89]
        x = np.linspace(-4, 4, 201)
        y = fits.GaussianFit._model(x, *p0)
        noise = np.random.randn(y.size) * 0.2
        y += noise
        fit = fits.GaussianFit(x, y, make_plots=False)
        self.assertFitInterval(p0[0], "B", fit)
        self.assertFitInterval(p0[1], "A", fit)
        self.assertFitInterval(p0[2], "μ", fit)
        self.assertFitInterval(p0[3], "σ", fit)

    def test_MultiGaussianFit(self):
        p = [0.35, 2.8, 2.04, 0.88, 1.93, -2.3, 1.19]
        x = np.linspace(-10, 10)
        y = fits.MultiGaussianFit._model(x, *p)
        noise = np.random.randn(y.size) * 0.2
        y += noise
        fit = fits.MultiGaussianFit(x, y, make_plots=False, n_gaussians=2)

        #Be careful since no guarantee of order of fits
        #also only testing means and std devs since the other parameters are still a
        #little flaky...
        if fit.fit_params["μ0"] < fit.fit_params["μ1"]:
            self.assertFitInterval(p[5], "μ0", fit)
            self.assertFitInterval(p[2], "μ1", fit)
            self.assertFitInterval(p[6], "σ0", fit)
            self.assertFitInterval(p[3], "σ1", fit)
        else:
            self.assertFitInterval(p[2], "μ0", fit)
            self.assertFitInterval(p[5], "μ1", fit)
            self.assertFitInterval(p[3], "σ0", fit)
            self.assertFitInterval(p[6], "σ1", fit)


    def test_T1Fit(self):

        p0 = [2.0, 15, -1]
        x = np.linspace(0, 80, 201)
        y = qubit_fits.T1Fit._model(x, *p0)
        noise = np.random.randn(y.size)/10
        y += noise
        fit = qubit_fits.T1Fit(x, y, make_plots=False)
        self.assertFitInterval(p0[0], "A", fit)
        self.assertFitInterval(p0[1], "T1", fit)
        self.assertFitInterval(p0[2], "A0", fit)

    def test_RabiAmpFit(self):

        p0 = [0.3, 0.76, 1.0, 0.01]
        x = np.linspace(-0.5, 0.5, 201)
        y = qubit_fits.RabiAmpFit._model(x, *p0)
        noise = np.random.randn(y.size) * np.max(y)/10
        y += noise
        fit = qubit_fits.RabiAmpFit(x, y, make_plots=False)
        self.assertFitInterval(p0[1], "Api", fit)

    def test_RabiWidthFit(self):

        p0 = [0.1, 0.34, 9, 0.2, 0.02]
        x = np.linspace(0, 20, 201)
        y = qubit_fits.RabiWidthFit._model(x, *p0)
        noise = np.random.randn(y.size) * np.max(y)/10
        y += noise
        fit = qubit_fits.RabiWidthFit(x, y, make_plots=False)
        self.assertFitInterval(p0[2], "T", fit)

    def test_RamseyFit_1f(self):
        #x, f, A, tau, phi, y0

        p0 = [0.22, 1.0, 11.3, 0.01, 0.1]
        x = np.linspace(0, 30, 201)
        y = qubit_fits.RamseyFit._model_1f(x, *p0)
        noise = np.random.randn(y.size) * np.max(y)/10
        y += noise

        fit = qubit_fits.RamseyFit(x, y, two_freqs=False, make_plots=False)
        self.assertFitInterval(p0[0], "f", fit)
        self.assertFitInterval(p0[2], "tau", fit)

    def test_RamseyFit_2f(self):
        #x, f, A, tau, phi, y0

        p0 = [0.22, 0.1, 1.0, 0.6, 11.3, 16.4, 0.01, 0.02, 0.1, 0.08]
        x = np.linspace(0, 30, 201)
        y = qubit_fits.RamseyFit._model_2f(x, *p0)
        noise = np.random.randn(y.size) * np.max(y)/10
        y += noise

        fit = qubit_fits.RamseyFit(x, y, two_freqs=True, AIC=True, make_plots=False, force=True)
        self.assertFitInterval(p0[0], "f1", fit)
        self.assertFitInterval(p0[1], "f2", fit)
        self.assertFitInterval(p0[4], "tau1", fit)
        self.assertFitInterval(p0[5], "tau2", fit)

    @unittest.skip("Fit not particularly stable?")
    def test_SingleQubitRB(self):
        p0 = [0.99, 5e-3, 0.2]

        x = np.array([2**n for n in range(10)])
        y = qubit_fits.SingleQubitRBFit._model(x, *p0)
        noise = np.random.randn(y.size) * p0[0]/100
        y += noise

        fit = qubit_fits.SingleQubitRBFit(x, y, make_plots=False)
        self.assertFitInterval(p0[1], "r", fit, tol=20)



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