import csv
import json
import logging
import pickle
from os import path
from pathlib import Path
from typing import Generator, Tuple, Optional, List, cast

import dill

from autoconf import cached_property
from autofit.non_linear.samples.pdf import SamplesPDF
from autofit.aggregator.file_output import (
    JSONOutput,
    FileOutput,
)
from autofit.mapper.identifier import Identifier
from autofit.non_linear.samples.sample import samples_from_iterator
from autoconf.dictable import from_dict

# noinspection PyProtectedMember
original_create_file_handle = dill._dill._create_filehandle


def _create_file_handle(*args, **kwargs):
    """
    Handle FileNotFoundError when attempting to deserialize pickles
    using dill and return None instead.
    """
    try:
        return original_create_file_handle(*args, **kwargs)
    except pickle.UnpicklingError as e:
        if not isinstance(e.args[0], FileNotFoundError):
            raise e
        logging.warning(
            f"Could not create a handler for {e.args[0].filename} as it does not exist"
        )
        return None


dill._dill._create_filehandle = _create_file_handle


class AbstractSearchOutput:
    def __init__(self, directory: Path, reference: Optional[dict] = None):
        self.directory = directory
        self._reference = reference

    @property
    def is_complete(self) -> bool:
        """
        Whether the search has completed
        """
        return (self.directory / ".completed").exists()

    @property
    def parent_identifier(self) -> Optional[str]:
        """
        Read the parent identifier for a fit in a directory.

        Defaults to None if no .parent_identifier file is found.
        """
        try:
            return (self.directory / ".parent_identifier").read_text()
        except FileNotFoundError:
            return None

    @property
    def files_path(self):
        return self.directory / "files"

    def _outputs(self, suffix):
        outputs = []
        for file_path in self.files_path.rglob(f"*{suffix}"):
            name = ".".join(
                file_path.relative_to(self.files_path).with_suffix("").parts
            )
            outputs.append(FileOutput(name, file_path))
        return outputs

    @cached_property
    def jsons(self) -> List[JSONOutput]:
        """
        The json files in the search output files directory
        """
        return cast(List[JSONOutput], self._outputs(".json"))

    @cached_property
    def arrays(self):
        """
        The csv files in the search output files directory
        """
        return self._outputs(".csv")

    @cached_property
    def pickles(self):
        """
        The pickle files in the search output files directory
        """
        return self._outputs(".pickle")

    @property
    def hdus(self):
        """
        The fits files in the search output files directory
        """
        return self._outputs(".fits")

    @property
    def max_log_likelihood(self) -> Optional[float]:
        """
        The log likelihood of the maximum log likelihood sample
        """
        try:
            return self.samples.max_log_likelihood_sample.log_likelihood
        except AttributeError:
            return None

    def __getattr__(self, name):
        """
        Attempt to load a pickle by the same name from the search output directory.

        dataset.pickle, meta_dataset.pickle etc.
        """
        return self.value(name)

    def value(self, name: str):
        """
        Load the value of some object in the files directory for the search.

        This may be a pickle, json, csv or fits file.

        If the JSON has a specified type it is parsed as that type. See dictable.py
        in autoconf.

        Returns None if the file does not exist.

        Parameters
        ----------
        name
            The name of the file to load without a file suffix.

        Returns
        -------
        The loaded object
        """
        for item in self.jsons:
            if item.name == name:
                return item.value_using_reference(self._reference)
        for item in self.pickles + self.arrays + self.hdus:
            if item.name == name:
                return item.value

        return None


class SearchOutput(AbstractSearchOutput):
    """
    @DynamicAttrs
    """

    is_grid_search = False

    def __init__(self, directory: Path, reference: dict = None):
        """
        Represents the output of a single search. Comprises a metadata file and other dataset files.

        Parameters
        ----------
        directory
            The directory of the search
        """
        super().__init__(directory, reference)
        self.__search = None
        self.__model = None
        self._samples = None

        self.directory = directory

        self.file_path = directory / "metadata"

        try:
            with open(self.file_path) as f:
                self.text = f.read()
                pairs = [
                    line.split("=") for line in self.text.split("\n") if "=" in line
                ]
                self.__dict__.update({pair[0]: pair[1] for pair in pairs})
        except FileNotFoundError:
            pass

    @property
    def instance(self):
        """
        The instance of the maximum log likelihood sample i.e. the instance
        with the greatest likelihood.

        None if samples cannot be loaded.
        """
        try:
            return self.samples.max_log_likelihood()
        except (AttributeError, NotImplementedError):
            return None

    @property
    def id(self) -> str:
        """
        The unique identifier of the search.

        This is used as a directory name and as a database identifier.
        """
        return str(Identifier([self.search, self.model, self.unique_tag]))

    @property
    def model(self):
        """
        The model used by the search
        """
        if self.__model is None:
            self.__model = self.value("model")
        return self.__model

    @property
    def samples(self) -> SamplesPDF:
        """
        The samples of the search, parsed from a CSV containing individual samples
        and a JSON containing metadata.
        """
        if not self._samples:
            try:
                info_json = JSONOutput(
                    "info", self.files_path / "samples_info.json"
                ).dict

                with open(self.files_path / "samples.csv") as f:
                    sample_list = samples_from_iterator(csv.reader(f))

                self._samples = SamplesPDF.from_list_info_and_model(
                    sample_list=sample_list,
                    samples_info=info_json,
                    model=self.model,
                )
            except FileNotFoundError:
                raise AttributeError("No samples found")
        return self._samples

    def names_and_paths(
        self,
        suffix: str,
    ) -> Generator[Tuple[str, Path], None, None]:
        """
        Get the names and paths of files with a given suffix.

        Parameters
        ----------
        suffix
            The suffix of the files to retrieve (e.g. ".json")

        Returns
        -------
        A generator of tuples of the form (name, path) where name is the path to the file
        joined by . without the suffix and path is the path to the file
        """
        for file in list(self.files_path.rglob(f"*{suffix}")):
            name = ".".join(file.relative_to(self.files_path).with_suffix("").parts)
            yield name, file

    @property
    def child_analyses(self):
        """
        A list of child analyses loaded from the analyses directory
        """
        return list(map(SearchOutput, Path(self.directory).glob("analyses/*")))

    @property
    def model_results(self) -> str:
        """
        Reads the model.results file
        """
        with open(self.directory / "model.results") as f:
            return f.read()

    @property
    def mask(self):
        """
        A pickled mask object
        """
        with open(self.files_path / "mask.pickle", "rb") as f:
            return dill.load(f)

    @property
    def header(self) -> str:
        """
        A header created by joining the search name
        """
        phase = self.phase or ""
        dataset_name = self.dataset_name or ""
        return path.join(phase, dataset_name)

    @property
    def search(self):
        """
        The search object that was used in this phase
        """
        if self.__search is None:
            try:
                with open(self.files_path / "search.json") as f:
                    self.__search = from_dict(json.load(f))
            except (FileNotFoundError, ModuleNotFoundError):
                try:
                    with open(self.files_path / "search.pickle", "rb") as f:
                        self.__search = pickle.load(f)
                except (FileNotFoundError, ModuleNotFoundError):
                    logging.warning("Could not load search")
        return self.__search

    def child_values(self, name):
        """
        Get the values of a given key for all children
        """
        return [getattr(child, name) for child in self.child_analyses]

    @property
    def path_prefix(self):
        return self.search.paths.path_prefix

    @property
    def name(self):
        """
        The name of the search
        """
        return self.search.name

    @property
    def unique_tag(self):
        """
        The unique tag of the search
        """
        return self.search.unique_tag

    def __str__(self):
        return self.text

    def __repr__(self):
        return "<PhaseOutput {}>".format(self)


class GridSearchOutput(AbstractSearchOutput):
    is_grid_search = True

    @property
    def unique_tag(self) -> str:
        """
        The unique tag of the grid search.
        """
        with open(self.directory / ".is_grid_search") as f:
            return f.read()

    @property
    def id(self) -> str:
        """
        Use the unique tag of the grid search as an identifier.
        """
        return self.unique_tag


class GridSearch:
    def __init__(
        self,
        grid_search_output: GridSearchOutput,
        children: List[SearchOutput],
    ):
        """
        Represents the output of a grid search. Comprises overall information from the grid search
        and output from each individual search.

        Parameters
        ----------
        grid_search_output
            The output of the grid search
        children
            The outputs of each individual search performed as part of the grid search
        """
        self.grid_search_output = grid_search_output
        self.children = children

    @property
    def best_fit(self) -> SearchOutput:
        """
        The output for the search in the grid search that had the greatest log likelihood
        """
        return max(self.children, key=lambda x: x.instance.log_likelihood)

    def __getattr__(self, item):
        return getattr(self.grid_search_output, item)
