# -------------------------------------------------------------------------------------------------
# scientific
import numpy as np
import pandas as pd
# -------------------------------------------------------------------------------------------------
# PyQuantum.TC
from PyQuantum.TC_sink.WaveFunction import WaveFunction
# -------------------------------------------------------------------------------------------------
# PyQuantum.Common
from PyQuantum.Common.Matrix import *
# -------------------------------------------------------------------------------------------------



class DensityMatrix(Matrix):
    # -----------------------------------------------------------------------------------------------------------------
    # ---------------------------------------------------- INIT -------------------------------------------------------
    # -----------------------------------------------------------------------------------------------------------------
    def __init__(self, wf):
        Assert(isinstance(wf, WaveFunction), "wf is not WaveFunction")

        super(DensityMatrix, self).__init__(
            m=wf.m, n=wf.m, dtype=np.complex128)

        wf_data = wf.data

        ro_data = wf_data.dot(wf_data.getH())

        Assert(np.shape(ro_data) == (self.m, self.n), "size mismatch")

        self.data = ro_data
        # self.data = np.matrix(ro_data, dtype=np.complex128)

        self.states = wf.states
        self.size = np.shape(self.data)[0]
    # -----------------------------------------------------------------------------------------------------------------
    # -----------------------------------------------------------------------------------------------------------------



    # -----------------------------------------------------------------------------------------------------------------
    # ---------------------------------------------------- IPRINT -----------------------------------------------------
    # -----------------------------------------------------------------------------------------------------------------
    def iprint(self):
        df = pd.DataFrame()

        for i in range(self.size):
            for j in range(self.size):
                df.loc[i, j] = self.data[i, j]

        df.index = df.columns = [str(v) for v in self.states.values()]

        self.df = df
    # -----------------------------------------------------------------------------------------------------------------
    # -----------------------------------------------------------------------------------------------------------------



    # -----------------------------------------------------------------------------------------------------------------
    # ---------------------------------------------------- NORMALIZE --------------------------------------------------
    # -----------------------------------------------------------------------------------------------------------------
    def normalize(self):
        self.data = (self.data+self.data.getH())/2.0
        self.data /= np.sum(np.abs(self.data.diagonal()))

    def energy(self, capacity, n_atoms, states_bin, diag_abs):
        energy = [0] * (capacity + n_atoms+1)

        for i in range(1, len(states_bin)):
            energy[i] += np.sum(diag_abs[states_bin[i]])
            # for j in states_bin[i]:
            #     energy[i] += diag_abs[j]
            energy[i] *= i

        # for i in range(1, len(states_bin)):

        return energy
    # -----------------------------------------------------------------------------------------------------------------
    # -----------------------------------------------------------------------------------------------------------------

# =====================================================================================================================
