# This code originates from https://github.com/lacroixle/ztfsmp/ by Leander
# Lacroix

import gzip
from collections.abc import Iterable

import numpy as np
import pandas as pd


class PolokaListFile:
    """A parser for poloka *.list files"""

    def __init__(self, header, df, df_desc=None, df_format=None, filename=None):
        if not df_desc:
            df_desc = {}

        self.header = header
        self.df = df
        self.df_desc = df_desc
        self.df_format = df_format
        self.filename = filename

    @classmethod
    def from_filename(cls, filename, delim_whitespace=True):
        with open(filename) as f:
            header, df, df_desc, df_format = cls._read_list_ext(
                f, delim_whitespace=delim_whitespace
            )

        return cls(header, df, df_desc, df_format, filename=filename)

    @classmethod
    def from_gzip(cls, filename, delim_whitespace=True):
        with gzip.open(filename, "rt") as f:
            header, df, df_desc, df_format = cls._read_list_ext(
                f, delim_whitespace=delim_whitespace
            )

        return cls(header, df, df_desc, df_format, filename=filename)

    def write_to(self, filename):
        filename = str(filename)  # from Path if needed
        with open(filename, "w") as f:
            if self.header:
                for key in self.header.keys():
                    if isinstance(self.header[key], Iterable) and not isinstance(
                        self.header[key], str
                    ):
                        f.write(
                            "@{} {}\n".format(
                                key.upper(), " ".join(map(str, self.header[key]))
                            )
                        )
                    else:
                        f.write(f"@{key.upper()} {str(self.header[key])}\n")

            for column in self.df.columns:
                if column in self.df_desc.keys():
                    f.write(f"#{column} : {self.df_desc[column]}\n")
                else:
                    f.write(f"#{column} :\n")

            if self.df_format is not None:
                f.write(f"# format {self.df_format}\n")

            f.write("# end\n")
            self.df.to_csv(f, sep=" ", index=False, header=False)

    @staticmethod
    def _parse_header_line(line):
        line = line.split()
        key = line[0][1:].lower()

        try:
            value = np.asarray(line[1:], dtype=int).tolist()
        except ValueError:  # not castable to int, try float
            try:
                value = np.asarray(line[1:], dtype=float).tolist()
            except ValueError:  # not castable as float, keep it as str
                value = line[:1]
        if len(value) == 1:
            value = value[0]

        return key, value

    @classmethod
    def _read_list_ext(cls, f, delim_whitespace=True):
        # Extract global @ parameters
        header = {}
        line = f.readline().strip()
        curline = 1
        while line[0] == "@":
            key, value = cls._parse_header_line(line)
            header[key] = value

            line = f.readline().strip()
            curline += 1

        # Extract dataframe
        # First, column names
        columns = []
        df_format = None
        df_desc = {}
        while line[0] == "#":
            curline += 1
            line = line[1:-1].strip()

            splitted = line.split(" ")
            if splitted[0].strip() == "end":
                break

            if splitted[0].strip() == "format":
                df_format = str(" ".join(line.split(" ")[1:])).strip()
                line = f.readline()
                continue

            splitted = line.split(":")
            column = splitted[0].strip()
            columns.append(column)
            if len(splitted) > 1:
                df_desc[column] = splitted[1].strip()

            line = f.readline()

        df = pd.read_csv(
            f,
            sep=r"\s+" if delim_whitespace else " ",
            names=columns,
            index_col=False,
            skipinitialspace=True,
        )

        return header, df, df_desc, df_format
