import numpy as np
import pandas as pd
from munch import Munch
from plaster.tools.schema.schema import Schema as s
from plaster.tools.schema import check
from plaster.tools.utils import utils
from plaster.tools.aaseq.aaseq import aa_str_to_list
from plaster.run.priors import Priors, ParamsAndPriors


class SimV1Params(ParamsAndPriors):
    """
    Simulations parameters is and ErrorModel + parameters for sim
    """

    defaults = Munch(
        n_pres=1,
        n_mocks=0,
        n_edmans=1,
        n_samples_train=5_000,
        n_samples_test=1_000,
        dyes=[],
        labels=[],
        random_seed=None,
        train_n_sample_multiplier=None,  # This does not appear to be used anywhere. tfb
        allow_train_test_to_be_identical=False,
        enable_ptm_labels=False,
        is_survey=False,
    )

    schema = s(
        s.is_kws_r(
            is_survey=s.is_bool(),
            priors_desc=Priors.priors_desc_schema,
            n_pres=s.is_int(bounds=(0, None)),
            n_mocks=s.is_int(bounds=(0, None)),
            n_edmans=s.is_int(bounds=(0, None)),
            n_samples_train=s.is_int(bounds=(1, None)),
            n_samples_test=s.is_int(bounds=(1, None)),
            dyes=s.is_list(
                elems=s.is_kws_r(dye_name=s.is_str(), channel_name=s.is_str())
            ),
            labels=s.is_list(
                elems=s.is_kws_r(
                    aa=s.is_str(),
                    dye_name=s.is_str(),
                    label_name=s.is_str(),
                    ptm_only=s.is_bool(required=False, noneable=True),
                )
            ),
            channels=s.is_dict(required=False),
            random_seed=s.is_int(required=False, noneable=True),
            allow_train_test_to_be_identical=s.is_bool(required=False, noneable=True),
            enable_ptm_labels=s.is_bool(required=False, noneable=True),
        )
    )

    # def copy(self):
    #     # REMOVE everything that _build_join_dfs put in
    #     utils.safe_del(self, "df")
    #     utils.safe_del(self, "by_channel")
    #     utils.safe_del(self, "ch_by_aa")
    #
    #     dst = utils.munch_deep_copy(self, klass_set={SimV1Params})
    #     dst.error_model = ErrorModel(**dst.error_model)
    #     assert isinstance(dst, SimV1Params)
    #     return dst

    def __init__(self, **kwargs):
        super().__init__(source="SimV1Params", **kwargs)
        self._setup_dfs()

    def validate(self):
        super().validate()

        all_dye_names = list(set([d.dye_name for d in self.dyes]))

        # No duplicate dye names
        self._validate(
            len(all_dye_names) == len(self.dyes), "The dye list contains a duplicate"
        )

        # No duplicate labels
        self._validate(
            len(list(set(utils.listi(self.labels, "aa")))) == len(self.labels),
            "There is a duplicate label",
        )

        # All labels have a legit dye name
        [
            self._validate(
                label.dye_name in all_dye_names,
                f"Label {label.label_name} does not have a valid matching dye_name",
            )
            for label in self.labels
        ]

        # Channel mappings
        mentioned_channels = {dye.channel_name: False for dye in self.dyes}
        if "channels" in self:
            # Validate that channel mapping is complete
            for channel_name, ch_i in self.channels.items():
                self._validate(
                    channel_name in mentioned_channels,
                    f"Channel name '{channel_name}' was not found in dyes",
                )
                mentioned_channels[channel_name] = True

            self._validate(
                all([mentioned for _, mentioned in mentioned_channels.items()]),
                "Not all channels in dyes were enumerated in channels",
            )
        else:
            # No channel mapping: assign them
            self["channels"] = {
                ch_name: i
                for i, ch_name in enumerate(sorted(mentioned_channels.keys()))
            }

    @property
    def n_cycles(self):
        return self.n_pres + self.n_mocks + self.n_edmans

    def channel_names(self):
        return sorted(list(set(utils.listi(self.dyes, "channel_name"))))

    def channel_i_by_name(self):
        channels = self.channel_names()
        return {
            channel_name: channel_i for channel_i, channel_name in enumerate(channels)
        }

    @property
    def n_channels(self):
        return len(self.channel_i_by_name().keys())

    @property
    def n_channels_and_cycles(self):
        return self.n_channels, self.n_cycles

    def _setup_dfs(self):
        """
        The error model contains information about the dyes and labels and other terms.
        Those error model parameters are wired together by names which are useful
        for reconciling calibrations.

        But here, these "by name" parameters are all put into a dataframe so that
        they can be indexed by integers.
        """
        dyes_df = pd.DataFrame(self.dyes)
        assert len(dyes_df) > 0

        labels_df = pd.DataFrame(self.labels)
        assert len(labels_df) > 0

        # LOOKUP dye priors
        dye_priors = []
        for dye in self.dyes:
            # SEARCH priors by dye name and if not found by channel
            p_non_fluorescent = self.priors.get_exact(
                f"p_non_fluorescent.{dye.dye_name}"
            )
            if p_non_fluorescent is None:
                p_non_fluorescent = self.priors.get(
                    f"p_non_fluorescent.ch_{dye.channel_name}"
                )

            dye_priors += [
                Munch(dye_name=dye.dye_name, p_non_fluorescent=p_non_fluorescent.prior,)
            ]

        dye_priors_df = pd.DataFrame(dye_priors)
        # dye_priors_df: (dye_name, p_non_fluorescent)

        dyes_df = utils.easy_join(dyes_df, dye_priors_df, "dye_name")
        # dyes_df: (dye_name, channel_name, p_non_fluorescent)

        # TODO: LOOKUP label priors
        #       (p_failure_to_bind_aa, p_failure_to_attach_to_dye)

        # LOOKUP channel priors
        ch_priors = pd.DataFrame(
            [
                dict(
                    channel_name=channel_name,
                    ch_i=ch_i,
                    bg_mu=self.priors.get(f"bg_mu.ch_{ch_i}").prior,
                    bg_sigma=self.priors.get(f"bg_sigma.ch_{ch_i}").prior,
                    gain_mu=self.priors.get(f"gain_mu.ch_{ch_i}").prior,
                    gain_sigma=self.priors.get(f"gain_sigma.ch_{ch_i}").prior,
                    row_k_sigma=self.priors.get(f"row_k_sigma.ch_{ch_i}").prior,
                    p_bleach=self.priors.get(f"p_bleach.ch_{ch_i}").prior,
                )
                for channel_name, ch_i in self.channels.items()
            ]
        )
        # ch_priors: (channel_name, ch_i, ...)

        self._channel__priors = (
            utils.easy_join(dyes_df, ch_priors, "channel_name")
            .drop(columns=["p_non_fluorescent"])
            .drop_duplicates()
            .reset_index()
        )
        # self._channel__priors: (
        #    'ch_i', 'channel_name', 'bg_mu', 'bg_sigma', 'dye_name',
        #    'gain_mu', 'gain_sigma', 'index', 'p_bleach', 'row_k_sigma',
        # )

        # SANITY check channel__priors
        group_by_ch = self._channel__priors.groupby("ch_i")
        for field in (
            "bg_mu",
            "bg_sigma",
            "gain_mu",
            "gain_sigma",
            "row_k_sigma",
        ):
            assert np.all(group_by_ch[field].nunique() == 1)
        assert "p_non_fluorescent" not in self._channel__priors.columns

        labels_dyes_df = utils.easy_join(labels_df, dyes_df, "dye_name")
        self._dye__label__priors = utils.easy_join(
            labels_dyes_df, ch_priors, "channel_name"
        ).reset_index(drop=True)

        # self._dye__label__priors: (
        #     'channel_name', 'dye_name', 'aa', 'label_name',
        #     'ptm_only', 'p_non_fluorescent', 'ch_i', 'bg_mu', 'bg_sigma',
        #     'gain_mu', 'gain_sigma', 'row_k_sigma', 'p_bleach'
        # )

        self._ch_by_aa = {
            row.aa: row.ch_i for row in self._dye__label__priors.itertuples()
        }

    def dye__label__priors(self):
        """
        DataFrame(
            'channel_name', 'dye_name', 'aa', 'label_name',
            'ptm_only', 'p_non_fluorescent', 'ch_i', 'bg_mu', 'bg_sigma',
            'gain_mu', 'gain_sigma', 'row_k_sigma', 'p_bleach'
        )
        """
        return self._dye__label__priors

    def channel__priors(self):
        """
        DataFrame(
            'ch_i', 'channel_name', 'bg_mu', 'bg_sigma', 'dye_name',
            'gain_mu', 'gain_sigma', 'index', 'p_bleach', 'row_k_sigma',
        )
        """
        return self._channel__priors

    def by_channel(self):
        return self._channel__priors.set_index("ch_i")

    def to_label_list(self):
        """Summarize labels like: ["DE", "C"]"""
        return [
            "".join(
                [label.aa for label in self.labels if label.dye_name == dye.dye_name]
            )
            for dye in self.dyes
        ]

    def to_label_str(self):
        """Summarize labels like: DE,C"""
        return ",".join(self.to_label_list())

    @classmethod
    def construct_from_aa_list(cls, aa_list, **kwargs):
        """
        This is a helper to generate channel when you have a list of aas.
        For example, two channels where ch0 is D&E and ch1 is Y.
        ["DE", "Y"].

        If you pass in an error model, it needs to match channels and labels.
        """

        check.list_or_tuple_t(aa_list, str)

        allowed_aa_mods = ["[", "]"]
        assert all(
            [
                (aa.isalpha() or aa in allowed_aa_mods)
                for aas in aa_list
                for aa in list(aas)
            ]
        )

        dyes = [
            Munch(dye_name=f"dye_{ch}", channel_name=f"ch_{ch}")
            for ch, _ in enumerate(aa_list)
        ]

        # Note the extra for loop because "DE" needs to be split into "D" & "E"
        # which is done by aa_str_to_list() - which also handles PTMs like S[p]
        labels = [
            Munch(
                aa=aa, dye_name=f"dye_{ch}", label_name=f"label_{ch}", ptm_only=False,
            )
            for ch, aas in enumerate(aa_list)
            for aa in aa_str_to_list(aas)
        ]

        return cls(dyes=dyes, labels=labels, **kwargs)

import numpy as np
import pandas as pd
from munch import Munch
from plaster.tools.schema.schema import Schema as s
from plaster.tools.schema import check
from plaster.tools.utils import utils
from plaster.tools.aaseq.aaseq import aa_str_to_list
from plaster.run.priors import Priors, ParamsAndPriors


class SimV1Params(ParamsAndPriors):
    """
    Simulations parameters is and ErrorModel + parameters for sim
    """

    defaults = Munch(
        n_pres=1,
        n_mocks=0,
        n_edmans=1,
        n_samples_train=5_000,
        n_samples_test=1_000,
        dyes=[],
        labels=[],
        random_seed=None,
        train_n_sample_multiplier=None,  # This does not appear to be used anywhere. tfb
        allow_train_test_to_be_identical=False,
        enable_ptm_labels=False,
        is_survey=False,
    )

    schema = s(
        s.is_kws_r(
            is_survey=s.is_bool(),
            priors_desc=Priors.priors_desc_schema,
            n_pres=s.is_int(bounds=(0, None)),
            n_mocks=s.is_int(bounds=(0, None)),
            n_edmans=s.is_int(bounds=(0, None)),
            n_samples_train=s.is_int(bounds=(1, None)),
            n_samples_test=s.is_int(bounds=(1, None)),
            dyes=s.is_list(
                elems=s.is_kws_r(dye_name=s.is_str(), channel_name=s.is_str())
            ),
            labels=s.is_list(
                elems=s.is_kws_r(
                    aa=s.is_str(),
                    dye_name=s.is_str(),
                    label_name=s.is_str(),
                    ptm_only=s.is_bool(required=False, noneable=True),
                )
            ),
            channels=s.is_dict(required=False),
            random_seed=s.is_int(required=False, noneable=True),
            allow_train_test_to_be_identical=s.is_bool(required=False, noneable=True),
            enable_ptm_labels=s.is_bool(required=False, noneable=True),
        )
    )

    # def copy(self):
    #     # REMOVE everything that _build_join_dfs put in
    #     utils.safe_del(self, "df")
    #     utils.safe_del(self, "by_channel")
    #     utils.safe_del(self, "ch_by_aa")
    #
    #     dst = utils.munch_deep_copy(self, klass_set={SimV1Params})
    #     dst.error_model = ErrorModel(**dst.error_model)
    #     assert isinstance(dst, SimV1Params)
    #     return dst

    def __init__(self, **kwargs):
        super().__init__(source="SimV1Params", **kwargs)
        self._setup_dfs()

    def validate(self):
        super().validate()

        all_dye_names = list(set([d.dye_name for d in self.dyes]))

        # No duplicate dye names
        self._validate(
            len(all_dye_names) == len(self.dyes), "The dye list contains a duplicate"
        )

        # No duplicate labels
        self._validate(
            len(list(set(utils.listi(self.labels, "aa")))) == len(self.labels),
            "There is a duplicate label",
        )

        # All labels have a legit dye name
        [
            self._validate(
                label.dye_name in all_dye_names,
                f"Label {label.label_name} does not have a valid matching dye_name",
            )
            for label in self.labels
        ]

        # Channel mappings
        mentioned_channels = {dye.channel_name: False for dye in self.dyes}
        if "channels" in self:
            # Validate that channel mapping is complete
            for channel_name, ch_i in self.channels.items():
                self._validate(
                    channel_name in mentioned_channels,
                    f"Channel name '{channel_name}' was not found in dyes",
                )
                mentioned_channels[channel_name] = True

            self._validate(
                all([mentioned for _, mentioned in mentioned_channels.items()]),
                "Not all channels in dyes were enumerated in channels",
            )
        else:
            # No channel mapping: assign them
            self["channels"] = {
                ch_name: i
                for i, ch_name in enumerate(sorted(mentioned_channels.keys()))
            }

    @property
    def n_cycles(self):
        return self.n_pres + self.n_mocks + self.n_edmans

    def channel_names(self):
        return sorted(list(set(utils.listi(self.dyes, "channel_name"))))

    def channel_i_by_name(self):
        channels = self.channel_names()
        return {
            channel_name: channel_i for channel_i, channel_name in enumerate(channels)
        }

    @property
    def n_channels(self):
        return len(self.channel_i_by_name().keys())

    @property
    def n_channels_and_cycles(self):
        return self.n_channels, self.n_cycles

    def _setup_dfs(self):
        """
        The error model contains information about the dyes and labels and other terms.
        Those error model parameters are wired together by names which are useful
        for reconciling calibrations.

        But here, these "by name" parameters are all put into a dataframe so that
        they can be indexed by integers.
        """
        dyes_df = pd.DataFrame(self.dyes)
        assert len(dyes_df) > 0

        labels_df = pd.DataFrame(self.labels)
        assert len(labels_df) > 0

        # LOOKUP dye priors
        dye_priors = []
        for dye in self.dyes:
            # SEARCH priors by dye name and if not found by channel
            p_non_fluorescent = self.priors.get_exact(
                f"p_non_fluorescent.{dye.dye_name}"
            )
            if p_non_fluorescent is None:
                p_non_fluorescent = self.priors.get(
                    f"p_non_fluorescent.ch_{dye.channel_name}"
                )

            dye_priors += [
                Munch(dye_name=dye.dye_name, p_non_fluorescent=p_non_fluorescent.prior,)
            ]

        dye_priors_df = pd.DataFrame(dye_priors)
        # dye_priors_df: (dye_name, p_non_fluorescent)

        dyes_df = utils.easy_join(dyes_df, dye_priors_df, "dye_name")
        # dyes_df: (dye_name, channel_name, p_non_fluorescent)

        # TODO: LOOKUP label priors
        #       (p_failure_to_bind_aa, p_failure_to_attach_to_dye)

        # LOOKUP channel priors
        ch_priors = pd.DataFrame(
            [
                dict(
                    channel_name=channel_name,
                    ch_i=ch_i,
                    bg_mu=self.priors.get(f"bg_mu.ch_{ch_i}").prior,
                    bg_sigma=self.priors.get(f"bg_sigma.ch_{ch_i}").prior,
                    gain_mu=self.priors.get(f"gain_mu.ch_{ch_i}").prior,
                    gain_sigma=self.priors.get(f"gain_sigma.ch_{ch_i}").prior,
                    row_k_sigma=self.priors.get(f"row_k_sigma.ch_{ch_i}").prior,
                    p_bleach=self.priors.get(f"p_bleach.ch_{ch_i}").prior,
                )
                for channel_name, ch_i in self.channels.items()
            ]
        )
        # ch_priors: (channel_name, ch_i, ...)

        self._channel__priors = (
            utils.easy_join(dyes_df, ch_priors, "channel_name")
            .drop(columns=["p_non_fluorescent"])
            .drop_duplicates()
            .reset_index()
        )
        # self._channel__priors: (
        #    'ch_i', 'channel_name', 'bg_mu', 'bg_sigma', 'dye_name',
        #    'gain_mu', 'gain_sigma', 'index', 'p_bleach', 'row_k_sigma',
        # )

        # SANITY check channel__priors
        group_by_ch = self._channel__priors.groupby("ch_i")
        for field in (
            "bg_mu",
            "bg_sigma",
            "gain_mu",
            "gain_sigma",
            "row_k_sigma",
        ):
            assert np.all(group_by_ch[field].nunique() == 1)
        assert "p_non_fluorescent" not in self._channel__priors.columns

        labels_dyes_df = utils.easy_join(labels_df, dyes_df, "dye_name")
        self._dye__label__priors = utils.easy_join(
            labels_dyes_df, ch_priors, "channel_name"
        ).reset_index(drop=True)

        # self._dye__label__priors: (
        #     'channel_name', 'dye_name', 'aa', 'label_name',
        #     'ptm_only', 'p_non_fluorescent', 'ch_i', 'bg_mu', 'bg_sigma',
        #     'gain_mu', 'gain_sigma', 'row_k_sigma', 'p_bleach'
        # )

        self._ch_by_aa = {
            row.aa: row.ch_i for row in self._dye__label__priors.itertuples()
        }

    def dye__label__priors(self):
        """
        DataFrame(
            'channel_name', 'dye_name', 'aa', 'label_name',
            'ptm_only', 'p_non_fluorescent', 'ch_i', 'bg_mu', 'bg_sigma',
            'gain_mu', 'gain_sigma', 'row_k_sigma', 'p_bleach'
        )
        """
        return self._dye__label__priors

    def channel__priors(self):
        """
        DataFrame(
            'ch_i', 'channel_name', 'bg_mu', 'bg_sigma', 'dye_name',
            'gain_mu', 'gain_sigma', 'index', 'p_bleach', 'row_k_sigma',
        )
        """
        return self._channel__priors

    def by_channel(self):
        return self._channel__priors.set_index("ch_i")

    def to_label_list(self):
        """Summarize labels like: ["DE", "C"]"""
        return [
            "".join(
                [label.aa for label in self.labels if label.dye_name == dye.dye_name]
            )
            for dye in self.dyes
        ]

    def to_label_str(self):
        """Summarize labels like: DE,C"""
        return ",".join(self.to_label_list())

    @classmethod
    def construct_from_aa_list(cls, aa_list, **kwargs):
        """
        This is a helper to generate channel when you have a list of aas.
        For example, two channels where ch0 is D&E and ch1 is Y.
        ["DE", "Y"].

        If you pass in an error model, it needs to match channels and labels.
        """

        check.list_or_tuple_t(aa_list, str)

        allowed_aa_mods = ["[", "]"]
        assert all(
            [
                (aa.isalpha() or aa in allowed_aa_mods)
                for aas in aa_list
                for aa in list(aas)
            ]
        )

        dyes = [
            Munch(dye_name=f"dye_{ch}", channel_name=f"ch_{ch}")
            for ch, _ in enumerate(aa_list)
        ]

        # Note the extra for loop because "DE" needs to be split into "D" & "E"
        # which is done by aa_str_to_list() - which also handles PTMs like S[p]
        labels = [
            Munch(
                aa=aa, dye_name=f"dye_{ch}", label_name=f"label_{ch}", ptm_only=False,
            )
            for ch, aas in enumerate(aa_list)
            for aa in aa_str_to_list(aas)
        ]

        return cls(dyes=dyes, labels=labels, **kwargs)

Simulations parameters is and ErrorModel + parameters for sim # This does not appear to be used anywhere. tfb # def copy(self): #     # REMOVE everything that _build_join_dfs put in #     utils.safe_del(self, "df") #     utils.safe_del(self, "by_channel") #     utils.safe_del(self, "ch_by_aa") # #     dst = utils.munch_deep_copy(self, klass_set={SimV1Params}) #     dst.error_model = ErrorModel(**dst.error_model) #     assert isinstance(dst, SimV1Params) #     return dst # No duplicate dye names # No duplicate labels # All labels have a legit dye name # Channel mappings # Validate that channel mapping is complete # No channel mapping: assign them The error model contains information about the dyes and labels and other terms.         Those error model parameters are wired together by names which are useful         for reconciling calibrations.          But here, these "by name" parameters are all put into a dataframe so that         they can be indexed by integers. # LOOKUP dye priors # SEARCH priors by dye name and if not found by channel # dye_priors_df: (dye_name, p_non_fluorescent) # dyes_df: (dye_name, channel_name, p_non_fluorescent) # TODO: LOOKUP label priors #       (p_failure_to_bind_aa, p_failure_to_attach_to_dye) # LOOKUP channel priors # ch_priors: (channel_name, ch_i, ...) # self._channel__priors: ( #    'ch_i', 'channel_name', 'bg_mu', 'bg_sigma', 'dye_name', #    'gain_mu', 'gain_sigma', 'index', 'p_bleach', 'row_k_sigma', # ) # SANITY check channel__priors # self._dye__label__priors: ( #     'channel_name', 'dye_name', 'aa', 'label_name', #     'ptm_only', 'p_non_fluorescent', 'ch_i', 'bg_mu', 'bg_sigma', #     'gain_mu', 'gain_sigma', 'row_k_sigma', 'p_bleach' # ) DataFrame(             'channel_name', 'dye_name', 'aa', 'label_name',             'ptm_only', 'p_non_fluorescent', 'ch_i', 'bg_mu', 'bg_sigma',             'gain_mu', 'gain_sigma', 'row_k_sigma', 'p_bleach'         ) DataFrame(             'ch_i', 'channel_name', 'bg_mu', 'bg_sigma', 'dye_name',             'gain_mu', 'gain_sigma', 'index', 'p_bleach', 'row_k_sigma',         ) Summarize labels like: ["DE", "C"] Summarize labels like: DE,C This is a helper to generate channel when you have a list of aas.         For example, two channels where ch0 is D&E and ch1 is Y.         ["DE", "Y"].          If you pass in an error model, it needs to match channels and labels. # Note the extra for loop because "DE" needs to be split into "D" & "E" # which is done by aa_str_to_list() - which also handles PTMs like S[p]