#
# Copyright The NOMAD Authors.
#
# This file is part of NOMAD. See https://nomad-lab.eu for further info.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Stuttgart analysis results file generated by APyT."""

# pylint: disable=duplicate-code

import re

import flatdict as fd
import numpy as np
import pandas as pd
import yaml

from ifes_apt_tc_data_modeling.utils.custom_logging import logger
from ifes_apt_tc_data_modeling.utils.pint_custom_unit_registry import ureg


class ReadStuttgartApytMetadataFileFormat:
    """Read metadata of an entry from the db.yaml-style database-in-a-file as used by APyT."""

    def __init__(self, file_path: str, verbose: bool = False):
        self.supported = False
        if not file_path.lower().endswith((".yml", ".yaml")):
            logger.warning(f"{file_path} is likely not an APyT database.yaml file")
            return
        self.supported = True
        self.verbose = verbose
        self.file_path = file_path

    def get_metadata(self, database_entry: int = 0):
        """Read metadata of a specific database entry."""
        if not self.supported:
            return
        try:
            with open(self.file_path, encoding="utf-8") as stream:
                database = yaml.safe_load(stream)
                if database_entry in database:
                    self.flat_metadata = fd.FlatDict(
                        database[database_entry], delimiter="/"
                    )
                    if self.verbose:
                        for key, val in self.flat_metadata.items():
                            logger.info(f"key: {key}, value: {val}")
                else:
                    logger.warning(
                        f"{self.file_path} does not contain an entry named {database_entry}"
                    )
        except (OSError, FileNotFoundError):
            logger.warning(f"{self.file_path} either FileNotFound or IOError !")


class ReadStuttgartApytSpectrumAlignFileFormat:
    """Read results from the text file that an `apyt_spectrum_align --no-sql 1` call generates."""

    def __init__(self, file_path: str, verbose: bool = False):
        self.supported = False
        if not file_path.lower().endswith(".txt"):
            logger.warning(
                f"{file_path} is likely not an APyT complete spectrum txt file"
            )
            return
        with open(file_path) as fp:
            first_line = fp.readline().strip().rstrip("\n")
            if first_line != "# m/q (amu/e)	  counts":
                logger.warning(f"{file_path} header is malformed")
                return
            pattern = re.compile(
                r"^"
                r"[+-]\d+\.\d{3}"
                r"[\t ]+"
                r"\d+"
                r"$"
            )
            for line in fp:
                if pattern.fullmatch(line.strip().rstrip("\n")):
                    continue
                else:
                    logger.warning(f"{file_path} line {line} is malformed")
                    return
        self.supported = True
        self.verbose = verbose
        self.file_path = file_path

    def get_complete_spectrum(self):
        """Read complete (binned) mass spectrum."""
        if not self.supported:
            return
        try:
            data_frame = pd.read_csv(self.file_path, sep=r"\s+", comment="#")
            number_of_bins = np.shape(data_frame)[0]
            histogram_bin_edges = np.zeros((number_of_bins,), np.float32)
            histogram_bin_edges[:] = data_frame.iloc[:, 0]
            histogram_bin_counts = np.zeros((number_of_bins,), np.uint32)
            histogram_bin_counts[:] = data_frame.iloc[:, 1]
            return (
                ureg.Quantity(histogram_bin_edges, ureg.dalton),
                ureg.Quantity(histogram_bin_counts),
            )
        except (OSError, FileNotFoundError):
            logger.warning(f"{self.file_path} either FileNotFound or IOError !")
            return


class ReadStuttgartApytSpectrumFitFileFormat:
    """Read results from text file(s) that an `apyt_spectrum_fit --no-sql 1 "{'W': ((3,), 0.0158)}` call generates."""

    def __init__(self, file_path: str):
        pass


class ReadStuttgartApytReconstructionFileFormat:
    """Read results from the text file that an `apyt_reconstruction --no-sql --module classic 1` call generates."""

    def __init__(self, file_path: str, verbose: bool = False):
        self.supported = False
        if not file_path.lower().endswith("_xyz.txt"):
            logger.warning(f"{file_path} is likely not an APyT _xyz.txt file")
            return
        with open(file_path) as fp:
            number_of_events = int(fp.readline().strip())
            skip = fp.readline().strip()
            pattern = re.compile(
                r"^"
                r"\d+"
                r"[\t ]+"
                r"[+-]?\d+\.\d+e[+-]\d{2}"
                r"[\t ]+"
                r"[+-]?\d+\.\d+e[+-]\d{2}"
                r"[\t ]+"
                r"[+-]?\d+\.\d+e[+-]\d{2}"
            )
            check_sum = 0
            for line in fp:
                if pattern.match(
                    line.strip().rstrip("\n")
                ):  # relax to not catch last column recon volume per atom
                    parts = line.split(maxsplit=1)
                    try:
                        ion_id = int(parts[0])
                        if 0 < ion_id <= 255:
                            check_sum += 1
                            continue
                        else:
                            logger.warning(
                                f"{file_path} has a line with an ion id that exceeds the supported value range"
                            )
                    except ValueError:
                        logger.warning(
                            f"{file_path} has a line with an ion id that is incorrectly formatted"
                        )
                        return
                else:
                    logger.warning(f"{file_path} line {line} is malformed")
                    return
            if number_of_events != check_sum:
                logger.warning(
                    f"{file_path} mismatch between number of events expected and found"
                )
                return
        self.supported = True
        self.verbose = verbose
        self.file_path = file_path
        self.number_of_events = number_of_events

    def get_reconstructed_positions(self):
        """Read xyz columns."""
        if not self.supported:
            return
        values = np.asarray(
            np.loadtxt(self.file_path, skiprows=2, usecols=(1, 2, 3)), np.float32
        )
        if np.shape(values)[0] != self.number_of_events:
            logger.warning(
                f"{self.file_path} number of reconstructed positions and events not matching"
            )
            return
        return ureg.Quantity(values, ureg.nanometer)

    def get_ion_type(self):
        """Read ion identity."""
        if not self.supported:
            return
        values = np.asarray(
            np.loadtxt(self.file_path, skiprows=2, usecols=(0)), np.uint8
        )
        if np.shape(values)[0] != self.number_of_events:
            logger.warning(
                f"{self.file_path} number of ion type ids and events not matching"
            )
            return
        return values
