"""
This file contains the parser of the density matrix generated by TB2J.
"""

import pickle
import numpy as np
from dataclasses import dataclass


@dataclass
class DensityMatrix:
    """
    Rlist: list of R vectors, np.ndarray of shape (NR, 3), or None
    rhoR: density matrix, np.ndarray of shape (nR, N, N), or None
    Rmap: dict of R vectors, or None
    """

    Rlist: np.ndarray
    rhoR: np.ndarray
    Rmap: dict

    def __init__(self, Rlist=None, rhoR=None):
        self.Rlist = Rlist
        self.rhoR = rhoR
        self._build_Rmap()

    def _build_Rmap(self):
        self.Rmap = {}
        for R in self.Rlist:
            self.Rmap[tuple(R)] = len(self.Rmap)

    def print_dm_at_R(self, R):
        """
        print the real part with 5 decimals of the density matrix at R
        """
        print(np.real(self.rhoR[self.Rmap[tuple(R)]]).round(5))

    def print_dm_at_R_block(self, R, ilist, jlist):
        """
        print the real part with 3 decimals of the density matrix at R, only the block of ilist and jlist
        """
        print("real part")
        print(np.real(self.rhoR[self.Rmap[tuple(R)]][np.ix_(ilist, jlist)]).round(5))
        print("imag part")
        print(np.imag(self.rhoR[self.Rmap[tuple(R)]][np.ix_(ilist, jlist)]).round(5))

    def get_dm_at_R(self, R):
        return self.rhoR[self.Rmap[tuple(R)]]

    @classmethod
    def load_pickle(self, fname):
        with open(fname, "rb") as f:
            d = pickle.load(f)
        return DensityMatrix(**d)
