"""
This module contains classes to parse the .dat files created by QucsStudio simulations.

Example usage:
    ds = DataSet("simulation.dat")
    plt.plot(ds.variables["frequency"], ds.variables["node1.V"])

For a more memory friendly implementation, check the DataSetHeader class.
"""
__author__ = "Leonhard Petzel"

from io import BufferedReader
import math
from typing import List, Self
import numpy as np


class DataSetVariableHeader(object):
    """ Represents one variable header entry within a qucs .dat file.

    Use DataSetVariableHeader.from_stream(file_handle) for initialization.
    File handle must be at the correct position.
    Use .read_from_file to obtain the content of the variable from a file as numpy array."""
    def __init__(self):
        "Use DataSetVariableHeader.from_stream to initialize."
        pass

    @classmethod
    def from_stream(cls, fh: BufferedReader, pos: int):
        """ Named constructor to parse a data set variable entry from the current file position.
        
        Arguments:
            fh: Binary file handle
            pos: position for reading the variable in the file. This should be 16 + header length + sum(variable.count_bytes for variable in variables_so_far). """
        return cls()._parse_bytes(fh, pos)

    def _parse_bytes(self, fh: BufferedReader, pos: int):
        read_int = lambda: int.from_bytes(fh.read(4), "little")
        read_str = lambda: b"".join(iter(lambda: fh.read(1), b"\0")).decode()
        bitget = lambda v, i: (v >> i) & 1

        self.position = pos

        self.type = read_int()
        self.complex = bitget(self.type, 2) == 0
        self.count = read_int()
        self.name = read_str()
        self.is_dependant = bitget(self.type, 1) == 1
        if self.is_dependant:
            self.dep_names = read_str().split(" ")
            self.dep_names.reverse()
        else:
            self.dep_names = []
        print(f"type={self.type:X}, len={self.count}, name={self.name}, deps={self.dep_names}, pos={self.position}")
        return self
    
    @property
    def count_bytes(self) -> int:
        """Number of bytes, this variable occupies in the file"""
        return self.count * (16 if self.complex else 8)
    
    def read_from_file(self, fh: BufferedReader):
        """ Takes the file descriptor of the qucs .dat file and parses the variable from it using the knowledge of position and length.
        The shape cannot be determined, because the dimensions of the dependencies are now known. The result is therefore flat. """
        fh.seek(self.position)
        cnt = fh.read(self.count_bytes)
        res = np.frombuffer(cnt, np.complex128 if self.complex else np.double) # TODO: check, if np.fromfile can be used
        return res


class DataSetHeader(object):
    """ Class representing the header section of a qucs .dat file.
    Use DataSetHeader.from_stream() to parse the header file.
    This class can be used, if you are interested in only some variables and don't want to load all.
    If you dont care about memory consumption, use DataSet() instead.

    Usage:
    
      with open(fname, "rb") as fh:
          header = DataSetHeader.from_stream(fh)
          voltage = header.read_variable(fh, 'node.V')
          freq = header.read_variable(fh, 'frequency')
      plt.plot(freq, volt)
    """
    def __init__(self) -> None:
        "Use the from_stream named constructor for initialization."
        pass

    @classmethod
    def from_stream(cls, fh: BufferedReader) -> "DataSetHeader":
        """ Parse the qucs .dat data set header.
        """
        return cls()._parse_bytes(fh)._correct_variables()


    def _parse_bytes(self, fh: BufferedReader) -> Self:
        read_int = lambda: int.from_bytes(fh.read(4), "little")

        code = fh.read(8)
        assert code == b"QucsData"
        self.version = read_int()
        self.size = read_int()
        # header = fh.read(self.size)
        # assert len(header) == self.size

        self.variable_list = []
        self.variable_list: List[DataSetVariableHeader]

        pos0 = self.size + 8 + 4 + 4
        pos = pos0
        while fh.tell() < pos0:
            variable = DataSetVariableHeader()._parse_bytes(fh, pos)
            self.variable_list.append(variable)
            pos += variable.count_bytes
        
        self.variables = {v.name: v for v in self.variable_list}
        return self
    
    def _correct_variables(self) -> Self:
        """Correct some issues in the .dat files regarding wrong dependencies.
        See http://qucsstudio.de/forums/topic/dependent-variables-information-in-dat-files-is-not-correct."""
        for variable in self.variable_list:
            deps = [self.variables[n] for n in variable.dep_names]
            exp_len = math.prod(d.count for d in deps)
            length = variable.count
            if len(deps) > 0 and exp_len != length:

                print(f"WARNING: field {variable.name}({', '.join(variable.dep_names)}) should have {'x'.join(str(d.count) for d in deps)}={exp_len} elements, but has {length} elements. Trying to fix this issue.")
                # Assume just one missing
                for dep_ind in range(len(deps)):
                    dep = deps[dep_ind]
                    if length * dep.count == exp_len:
                        print(f"Going to discard dependency {dep.name} from {variable.name} to achieve correct shape.")
                        del variable.dep_names[dep_ind]
                        break
                else:
                    print(f"Error: could not resolve conflict")
        return self

    
    def read_variable(self, fh: BufferedReader, name: str):
        """ Read a variable from the dataset stored in a .dat file.
        
        Args:
            fh: Binary file handle, from which to read the variable from.
            name: Variable name.
        
        Returns:
            np.ndarray: Array of double or complex. The shape depends on the dependent variables. """
        var = self.variables[name]
        shape = tuple(self.variables[dep].count for dep in var.dep_names)
        res = var.read_from_file(fh)
        if len(shape) > 1:
            res = res.reshape(shape)
        return res

class DataSet(object):
    """ Complete DataSet representing a qucs .dat file.
    
    All variables from the file are immediatly parsed and are accessible as .variables.
    The dependencies are accassible as .dependencies and .header contains more detailed header information.
    The dependencies are essential for index manipulations and the shape of the variables.
    # TODO: consider to provide Panda access.
    
    If memory consumption is of concern, the DataSetHeader class can be used directly to just read the variables of interest."""
    def __init__(self, fname: str) -> None:
        with open(fname, "rb") as fh:
            self.header = DataSetHeader.from_stream(fh)
            self.variables = {name: self.header.read_variable(fh, name) for name in self.header.variables.keys()}
            self.dependencies = {v.name: v.dep_names for v in self.header.variable_list}


if __name__ == "__main__":
    import sys
    from matplotlib import pyplot as plt
    ds = DataSet(sys.argv[1])
    print(ds.variables.keys())
    print(ds.variables["Pr1.Ib"])
    # x = ds.variables[ds.dependencies["Pr1.Ib"][1]]
    # y0 = ds.variables["Pr1.Ib"][1,:]
    # y1 = ds.variables["Pr1.Ib"][2,:]
    # y2 = ds.variables["Pr1.Ib"][3,:]
    x = ds.variables[ds.dependencies["Pr1.Ib"][0]]
    y0 = ds.variables["Pr1.Ib"][:,1]
    y1 = ds.variables["Pr1.Ib"][:,2]
    y2 = ds.variables["Pr1.Ib"][:,3]
    plt.plot(x, 10*np.log10(np.abs(y0**2*50))+30, label="Fundamental")
    plt.plot(x, 10*np.log10(np.abs(y1**2*50))+30, label="2nd")
    plt.plot(x, 10*np.log10(np.abs(y2**2*50))+30, label="3rd")
    plt.legend()
    plt.show()