'''
Toolset to analyse measurement-induced Dephasing of qubits

Hacked together by Rene Vollmer
'''

import pycqed
from pycqed.analysis_v2.quantum_efficiency_analysis import RamseyAnalysisSweep
import pycqed.analysis_v2.base_analysis as ba
import numpy as np
from collections import OrderedDict
import copy
import datetime
import os
from pycqed.analysis import analysis_toolbox as a_tools

import numpy as np
import matplotlib.pyplot as plt
import numpy.ma
from mpl_toolkits.axes_grid1.inset_locator import inset_axes


class CrossDephasingAnalysis(ba.BaseDataAnalysis):
    '''
    Analyses measurement-induced Dephasing of qubits
    '''

    def __init__(self, qubit_labels: list,
    			 t_start: str = None, t_stop: str = None,
                 label_pattern: str = 'ro_amp_sweep_ramsey_trgt_{TQ}_measured_{RQ}',
                 options_dict: dict = None,
                 extract_only: bool = False, auto: bool = True,
                 close_figs: bool = True, do_fitting: bool = True):

        super().__init__(t_start=t_start, t_stop=t_stop,
                         label=label_pattern,
                         options_dict=options_dict,
                         do_fitting=do_fitting,
                         close_figs=close_figs,
                         extract_only=extract_only)

        self.label_pattern = label_pattern
        self.qubit_labels = qubit_labels
        self.ra = np.array([[None] * len(qubit_labels)] * len(qubit_labels))
        d = copy.deepcopy(self.options_dict)
        d['save_figs'] = False
        for i, tq in enumerate(qubit_labels):
            for j, rq in enumerate(qubit_labels):
                label = label_pattern.replace('{TQ}', tq).replace('{RQ}', rq)
                self.ra[i, j] = RamseyAnalysisSweep(
                                                t_start=t_start,
                                                t_stop=t_stop,
                                                label=label, options_dict=d,
                                                auto=False, extract_only=True)

        if auto:
            self.run_analysis()

    def extract_data(self):
        ts = []
        for i, tq in enumerate(self.qubit_labels):
            for j, rq in enumerate(self.qubit_labels):
                ra = self.ra[i, j]
                ra.extract_data()
                ts.append(np.max(ra.raw_data_dict['datetime']))

        youngest = np.max(ts)
        youngest += datetime.timedelta(seconds=1)

        self.raw_data_dict = OrderedDict()
        self.raw_data_dict['datetime'] = [youngest]
        self.raw_data_dict['timestamps'] = [youngest.strftime("%Y%m%d_%H%M%S")]
        self.timestamps = [youngest.strftime("%Y%m%d_%H%M%S")]

        f = '%s_measurement_cross_dephasing_analysis' % (youngest.strftime("%H%M%S"))
        d = '%s' % (youngest.strftime("%Y%m%d"))
        folder = os.path.join(a_tools.datadir, d, f)
        self.raw_data_dict['folder'] = [folder]
        self.options_dict['analysis_result_file'] = os.path.join(folder, f + '.hdf5')

    def run_fitting(self):
        qubit_labels = self.qubit_labels
        self.fit_dicts = OrderedDict()
        self.fit_res = OrderedDict()
        self.fit_dicts['sigmas'] = np.array(
            [[None] * len(qubit_labels)] * len(qubit_labels), dtype=float)
        self.fit_dicts['sigmas_norm'] = np.array(
            [[None] * len(qubit_labels)] * len(qubit_labels), dtype=float)
        self.fit_dicts['deph_norm'] = np.array(
            [[None] * len(qubit_labels)] * len(qubit_labels), dtype=float)
        self.fit_res['coherence_fit'] = np.array(
            [[None] * len(qubit_labels)] * len(qubit_labels), dtype=object)

        for i, tq in enumerate(qubit_labels):
            for j, rq in enumerate(qubit_labels):
                ra = self.ra[i, j]
                ra.run_analysis()
                self.fit_res['coherence_fit'] = ra.fit_res['coherence_fit']
                self.fit_dicts['sigmas'][i, j] = ra.fit_dicts['coherence_fit']['sigma']
            self.fit_dicts['sigmas_norm'][i,:] = self.fit_dicts['sigmas'][i,:] / self.fit_dicts['sigmas'][i, i]
            self.fit_dicts['deph_norm'][i,:] = self.fit_dicts['sigmas'][i, i] / self.fit_dicts['sigmas'][i,:]

    def prepare_plots(self):
        self.plot_dicts['sigmas'] = {
            'plotfn': self.plot_labeled_2d,
            'title': '',  # todo