from .parsers import *
from .util import (
    add_braces_if_no_braces,
    dot2int,
    int2dot,
    seq2int,
    standardize_sequence,
    sequence_has_regular_characters,
)
import numpy as np
from .rnastructure import RNAstructure_singleton
from .env import env


class Datapoint:

    """A datapoint is a data structure where:
    - reference is the name of the sequence in the fasta file
    - sequence is a string of A, C, G, U
    - structure is a string of (, ), .
    - dms is a list of floats corresponding to the reactivity of each base in the sequence
    - paired_bases is a list of tuples (i,j) where i and j are paired bases.


    Example:
    >>> print(Datapoint(reference='reference', sequence='AACCGG', structure='((..))', dms=[1.0, 2.0, 3.0]))
    "reference":{"sequence": "AACCGG", "paired_bases": [[0, 5], [1, 4]], "dms": [1.0, 2.0, 3.0]}
    >>> print(Datapoint(reference='reference', sequence='not a valid sequence', structure='((..))'))
    None
    >>> print(Datapoint(reference='reference', sequence='NNNNNNN', structure='((..))'))
    None
    >>> print(Datapoint(reference='reference', sequence='AACCGG', paired_bases=[(0, 5), (1, 4)]))
    "reference":{"sequence": "AACCGG", "paired_bases": [[0, 5], [1, 4]]}
    >>> print(Datapoint(reference='reference', sequence='AACCGG', paired_bases=[(1, 5), (1, 4)]))
    None
    """

    def __new__(cls, *args, **kwargs):
        # Check if the conditions are met before creating the object
        sequence = kwargs.get("sequence", args[0] if len(args) > 0 else None)
        reference = kwargs.get("reference", args[1] if len(args) > 1 else None)
        if sequence is None or reference is None:
            return None
        sequence = standardize_sequence(sequence)

        if (
            not sequence_has_regular_characters(sequence)
            or len(sequence) == 0
            or len(reference) == 0
        ):
            return None

        if "paired_bases" in kwargs:
            if not cls._assert_paired_bases(None, kwargs["paired_bases"], sequence):
                return None

        # If conditions are met, create and return the new instance
        instance = super().__new__(cls)
        return instance

    def __init__(
        self, sequence, reference, structure=None, dms=None, paired_bases=None
    ):
        for attr in [sequence, reference]:
            assert isinstance(
                attr, str
            ), f"Expected {attr} to be a string, got {type(attr)} instead."
            assert len(attr) > 0, f"Expected {attr} to be non-empty."

        # standardize the sequence
        sequence = standardize_sequence(sequence)

        if not sequence_has_regular_characters(sequence):
            raise Exception(
                f"Sequence {sequence} contains characters other than ACGTUacgtu."
            )

        self.reference = reference
        self.sequence = sequence
        self.structure = structure
        if paired_bases is not None or structure is not None:
            self.paired_bases = (
                self._format_paired_bases(paired_bases)
                if paired_bases is not None
                else self._format_paired_bases(
                    self.structure_to_paired_bases(structure)
                )
            )
        if dms is not None:
            self.dms = self._format_dms(dms)
        self.opt_dict = {
            attr: eval(f"self.{attr}")
            for attr in ["paired_bases", "dms"]
            if hasattr(self, attr)
        }

    def _format_paired_bases(self, paired_bases):
        """Returns a set of tuples.
        >>> datapoint = Datapoint(reference='reference', sequence='AACCGG', structure='((..))', dms=[1.0, 2.0, 3.0])
        >>> datapoint._format_paired_bases([[1, 4], [0, 5]])
        ((0, 5), (1, 4))
        """
        return (
            tuple(sorted([tuple(sorted(pair)) for pair in paired_bases]))
            if paired_bases is not None
            else None
        )

    def _assert_paired_bases(self, paired_bases, sequence):
        """Ensures that the base pairs are 0=<bp<len(sequence) and that that there's maximum one pair per base.
        >>> datapoint = Datapoint(reference='reference', sequence='AACCGG', structure='((..))', dms=[1.0, 2.0, 3.0])
        >>> datapoint._assert_paired_bases([(0, 5), (1, 4)], 'AACCGG')
        True
        >>> datapoint._assert_paired_bases([(0, 5), (1, 4), (1, 2)], 'AACCGG')
        False
        >>> datapoint._assert_paired_bases([(0, 5), (4, 1), (1, 2)], 'AACCGG')
        False
        >>> datapoint._assert_paired_bases([(0, 8), (3, 4), (1, 2)], 'AACCGG')
        False
        >>> datapoint._assert_paired_bases(None, 'AACCGG')
        True
        """
        if paired_bases is not None:
            return (
                all(
                    [
                        0 <= pair[0] < len(sequence)
                        and 0 <= pair[1] < len(sequence)
                        and pair[0] != pair[1]
                        for pair in paired_bases
                    ]
                )
                and len(
                    set(
                        [pair[0] for pair in paired_bases]
                        + [pair[1] for pair in paired_bases]
                    )
                )
                == len(paired_bases) * 2
            )
        else:
            return True

    def _format_dms(self, dms):
        """returns a tuple
        >>> datapoint = Datapoint(reference='reference', sequence='AACCGG', structure='((..))', dms=[1.0, 2.0, 3.0])
        >>> datapoint._format_dms([1.0, 2.0, 3.0])
        (1.0, 2.0, 3.0)
        """
        if type(dms) == np.ndarray:
            dms = dms.tolist()
        return tuple(dms)
    
    def embed_dms(self):
        """Returns a list of floats corresponding to the dms.

        >>> from numpy import array, float32
        >>> datapoint = Datapoint(reference='reference', sequence='AACCGG', structure='((..))', dms=[1.0, 2.0, 3.0])
        >>> assert (datapoint.embed_dms() == array([1., 2., 3.], dtype=float32)).all(), 'The dms are not embedded correctly.'
        """
        return np.array([round(d,4) for d in self.dms], dtype=np.float32)

    def to_dict(self):
        return self.reference, {
            "sequence": self.sequence,
            **{k: v for k, v in self.opt_dict.items() if v is not None},
        }

    def to_flat_dict(self):
        return {
            "reference": self.reference,
            "sequence": self.sequence,
            **{k: v for k, v in self.opt_dict.items() if v is not None},
        }

    @classmethod
    def from_flat_dict(cls, d):
        return cls(
            sequence=d["sequence"],
            reference=d["reference"],
            **{k: v for k, v in d.items() if k not in ["sequence", "reference"]},
        )

    def __str__(self):
        return (
            '"'
            + self.reference
            + '"'
            + ":"
            + str(
                {
                    "sequence": self.sequence,
                    **{k: v for k, v in self.opt_dict.items() if v is not None},
                }
            )
            .replace("'", '"')
            .replace("(", "[")
            .replace(")", "]")
            .replace("None", "null")
        )

    def __repr__(self) -> str:
        out = (
            f"{self.__class__.__name__}('{self.reference}', sequence='{self.sequence}'"
        )
        if hasattr(self, "paired_bases"):
            out += f", paired_bases={self.paired_bases}"
        if hasattr(self, "dms"):
            out += f", dms={self.dms}"
        return out + ")"

    def structure_to_paired_bases(self, structure):
        """Returns a list of tuples (i,j) where i and j are paired bases."""
        paired_bases = set()
        stack = []
        for i, char in enumerate(structure):
            if char == "(":
                stack.append(i)
            elif char == ")":
                paired_bases.add((stack.pop(), i))
        return paired_bases

    def embed_sequence(self):
        """Returns a list of integers corresponding to the sequence.

        >>> from numpy import array
        >>> datapoint = Datapoint(reference='reference', sequence='AACCGG', structure='((..))', dms=[1.0, 2.0, 3.0])
        >>> assert (datapoint.embed_sequence() == array([1, 1, 2, 2, 3, 3])).all(), 'The sequence is not embedded correctly.'
        """
        return np.array([seq2int[base] for base in self.sequence], dtype=np.int8)

    def embed_structure(self):
        """Returns a list of integers corresponding to the structure.

        >>> from numpy import array
        >>> datapoint = Datapoint(reference='reference', sequence='AACCGG', structure='((..))', dms=[1.0, 2.0, 3.0])
        >>> assert not (datapoint.embed_structure() - array([2, 2, 1, 1, 3, 3])).any(), 'The sequence is not embedded correctly.'
        """
        return np.array([dot2int[base] for base in self.structure], dtype=np.int8)


class DatapointFactory:
    """A datapoint is a tuple (sequence, reference, structure, dms) where:

    - reference is the name of the sequence in the fasta file
    - sequence is a string of A, C, G, U
    - structure is a string of (, ), .
    - dms is a list of floats corresponding to the reactivity of each base in the sequence

    """

    def from_ct(ct_file, predict_dms):
        """Create a datapoint from a ct file. If predict_dms is True, the dms will be predicted using RNAstructure"""
        reference, sequence, paired_bases = Ct.parse(ct_file)
        sequence = standardize_sequence(sequence)

        if sequence_has_regular_characters(sequence):
            return Datapoint(
                sequence,
                reference,
                paired_bases=paired_bases,
                dms=Ct.predict_dms(ct_file) if predict_dms else None,
            )

    def from_fasta(sequence, reference, predict_structure, predict_dms):
        """Create a datapoint from a fasta file. The structure and dms will be None."""
        sequence = standardize_sequence(sequence)

        if sequence_has_regular_characters(sequence):
            structure, dms = None, None
            if predict_structure:
                structure = Fasta.predict_structure(sequence)
            if predict_dms:
                dms = Fasta.predict_dms(sequence)
            return Datapoint(sequence, reference, structure, dms)

    def from_dreem_output(reference, sequence, mutation_rate, predict_structure):
        """Create a datapoint from a dreem output file. The structure and dms will be predicted if predict_structure and predict_dms are True."""
        sequence = standardize_sequence(sequence)
        mutation_rate = np.array([float(m) for m in mutation_rate], dtype=np.float32)
        if sequence_has_regular_characters(sequence):
            return Datapoint(
                sequence=sequence,
                reference=reference,
                structure=RNAstructure_singleton.predictStructure(
                    sequence, dms=mutation_rate
                )
                if predict_structure
                else None,
                dms=mutation_rate,
            )

    def from_json_line(ref, d, predict_structure=False, predict_dms=False):
        """Create a datapoint from a json line. The json line should have the following format:
        "reference": {"sequence": "sequence", "paired_bases": [[1, 2], [3,4]], "dms": [1.0, 2.0, 3.0]}

        Example:
        >>> DatapointFactory.from_json_line("reference", {"sequence": "ACAAGU"})
        Datapoint('reference', sequence='ACAAGU')
        >>> DatapointFactory.from_json_line("reference", {"sequence": "ACAAGU", "paired_bases": [[1, 2], [3, 4]], "dms": [1.0, 2.0, 3.0]})
        Datapoint('reference', sequence='ACAAGU', paired_bases=((1, 2), (3, 4)), dms=(1.0, 2.0, 3.0))
        >>> DatapointFactory.from_json_line("reference", {"sequence": "something else than ACGTUacgtu", "paired_bases": [[1, 2], [3, 4]], "dms": [1.0, 2.0, 3.0]})
        """

        # create the datapoint
        sequence = d["sequence"]
        sequence = standardize_sequence(sequence)

        if predict_structure and (not "structure" in d) and (not "paired_bases" in d):
            d["structure"] = RNAstructure_singleton.predictStructure(sequence)

        if predict_dms and (not "dms" in d):
            d["dms"] = RNAstructure_singleton.predictPairingProbability(sequence)

        if sequence_has_regular_characters(sequence):
            return Datapoint(
                reference=ref,
                sequence=sequence,
                structure=d["structure"] if "structure" in d else None,
                paired_bases=d["paired_bases"] if "paired_bases" in d else None,
                dms=d["dms"] if "dms" in d else None,
            )
