#
# Copyright (C) 2021 Kevin Thornton <krthornt@uci.edu>
#
# This file is part of fwdpy11.
#
# fwdpy11 is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# fwdpy11 is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with fwdpy11.  If not, see <http://www.gnu.org/licenses/>.
#
import json
import typing
import warnings

import numpy as np
import tskit

from .._fwdpy11 import Mutation
from .metadata import (DiploidMetadata, decode_individual_metadata,
                       decode_mutation_metadata)


class WrappedTreeSequence(object):
    """
    Encapsulates a :class:`tskit.TreeSequence` as `self.ts`.

    Instance methods give additional functionality that usually
    requires interacting with the metadata.

    Instances of this class may be created using a :class:`tskit.TreeSequence`
    or a call to :func:`fwdpy11.tskit_tools.load`.

    .. versionadded:: 0.15.0

    .. deprecated:: 0.17.0
    """

    def _toplevel_metadata_value(self, name):
        from fwdpy11.tskit_tools import get_toplevel_metadata

        return get_toplevel_metadata(self._ts, name)

    def __init__(self, ts: tskit.TreeSequence):
        warnings.warn(
            FutureWarning(
                "fwdpy11.tskit_tools.WrappedTreeSequence is deprecated. "
                "Please use tskit.load and standalone functions "
                "from fwdpy11.tskit_tools instead"
            )
        )
        found = False
        for row in ts.provenances():
            record = json.loads(row.record)
            if "software" in record:
                if record["software"]["name"] == "fwdpy11":
                    found = True
        if not found:
            raise ValueError("this tree sequence was not generated by fwdpy11")

        self._ts = ts

    def timepoints_with_individuals(self, *, decode_metadata=False):
        from fwdpy11.tskit_tools import iterate_timepoints_with_individuals

        warnings.warn(
            FutureWarning(
                "use fwdpy11.tskit_tools.iterate_timepoints_with_individuals instead"
            )
        )

        return iterate_timepoints_with_individuals(
            self._ts, decode_metadata=decode_metadata
        )

    def decode_individual_metadata(
        self, rows: typing.Optional[typing.Union[int, slice]] = None
    ) -> typing.List[DiploidMetadata]:
        """
        Decode individual metadata.

        See :func:`fwdpy11.tskit_tools.decode_individual_metadata` for details.
        """
        return decode_individual_metadata(self._ts, rows)

    def decode_mutation_metadata(
        self, rows: typing.Optional[typing.Union[int, slice]] = None
    ) -> typing.List[typing.Optional[Mutation]]:
        """
        Decode mutation metadata.

        See :func:`fwdpy11.tskit_tools.decode_mutation_metadata` for details.
        """
        return decode_mutation_metadata(self._ts, rows)

    @property
    def generation(self):
        return self._toplevel_metadata_value("generation")

    @property
    def seed(self):
        return self._toplevel_metadata_value("seed")

    @property
    def demes_graph(self):
        rv = self._toplevel_metadata_value("demes_graph")
        if rv is not None:
            import demes

            rv = demes.Graph.fromdict(rv)

        return rv

    @property
    def model_params(self):
        rv = self._toplevel_metadata_value("model_params")
        if rv is not None:
            import fwdpy11

            try:
                temp = {}
                for key, value in rv.items():
                    temp[key] = fwdpy11.ModelParams(**eval(value))
                rv = temp
            except:
                rv = fwdpy11.ModelParams(**eval(rv))

        return rv

    @property
    def data(self):
        rv = self._toplevel_metadata_value("data")
        if rv is None:
            return rv

        return rv

    @property
    def ts(self):
        return self._ts
