"""Module containing classes and functions related to the different assay/designs used by pixelator.

Copyright © 2022 Pixelgen Technologies AB.
"""

import enum
import json
from functools import cache
from typing import Any, List, Mapping, Optional, Set, Tuple

try:
    from typing import Self
except ImportError:
    from typing_extensions import Self

import pydantic
from pydantic import BaseModel

from pixelator.common.config.utils import load_yaml_file
from pixelator.common.types import PathType


class RegionType(str, enum.Enum):
    """Enum class representing different sequence types."""

    UMI = "umi"
    UPI = "upi"
    PBS = "pbs"
    BARCODE = "barcode"
    AMPLICON = "amplicon"


class SequenceType(str, enum.Enum):
    """Enum class representing different sequence input types.

    JOINED: The sequence given must be empty.
            The sequence will be generated by joining the sequences of the
            sub-regions in the order specified in the assay specification.

    FIXED: The sequence of the regions is fixed and given.
    RANDOM: The sequence of the regions is random
    """

    JOINED = "joined"
    FIXED = "fixed"
    RANDOM = "random"


DNA_CHARS = {"A", "C", "G", "T"}


class RegionModel(BaseModel):
    """Validation model for region configuration."""

    region_id: str
    region_type: RegionType
    name: str
    sequence_type: SequenceType
    sequence: Optional[str] = None
    min_len: Optional[int] = None
    max_len: Optional[int] = None
    data: Optional[Mapping[str, Any]] = {}
    regions: Optional[List["RegionModel"]] = []

    @pydantic.model_validator(mode="after")
    def check_valid_dna_string(self) -> Self:
        """Validate DNA strings.

        :return: the validated model
        :raises ValueError: if the string contains non DNA characters
        """
        if self.region_type is RegionType.PBS:
            is_valid = True
            if self.sequence:
                is_valid = all(character in DNA_CHARS for character in self.sequence)
            if not is_valid:
                raise ValueError("Not a valid DNA string: found non DNA characters")

        return self


class AssayModel(BaseModel):
    """Validation model for assay configuration.

    :ivar name: Name of the assay
    :ivar assay_spec: List of assay regions
    """

    # TODO: Add more metadata fields
    name: str
    assay_spec: List[RegionModel]


class Region:
    """Class representing a region in an assay.

    :ivar region_id: unique ID of the region
    :ivar region_type: type of the region
    :ivar name: name of the region
    :ivar sequence_type: type of the sequence
    :ivar sequence: string representation of the region
    :ivar min_len: minimum length of the region
    :ivar max_len: maximum length of the region
    :ivar regions: list of sub-regions
    :ivar data: additional user-defined region properties
    """

    def __init__(
        self,
        region_id: str,
        region_type: RegionType,
        name: str,
        sequence_type: SequenceType,
        sequence: str = "",
        min_len: Optional[int] = None,
        max_len: Optional[int] = None,
        regions: Optional[List["Region"]] = None,
        data: Optional[Mapping[str, Any]] = None,
    ) -> None:
        """Initialize a Region."""
        self.parent_id: str | None = None
        self.region_id = region_id
        self.region_type = region_type
        self.name = name
        self.sequence_type = sequence_type
        self.sequence = sequence
        self.data = data

        self.min_len: int
        self.max_len: int

        if self.sequence_type == SequenceType.FIXED:
            if min_len != max_len or min_len != len(self.sequence):
                raise ValueError("Exact sequence does not match min/max length")

            if min_len is None:
                self.min_len = len(self.sequence)

            if max_len is None:
                self.max_len = len(self.sequence)

        self.min_len = min_len if min_len is not None else 0
        self.max_len = max_len if max_len is not None else 0
        self.regions = regions

        if self.regions:
            map(lambda r: r.set_parent_id(self.region_id), self.regions)

            self.min_len, self.max_len = self.get_len()
            self.sequence = self.get_sequence()

    def set_parent_id(self, parent_id: str) -> None:
        """Set the parent id of this region.

        :param parent_id: parent id to set
        """
        self.parent_id = parent_id

    def get_sequence(self) -> str:
        """Return the sequence representation of this region."""
        s = ""
        if self.regions:
            for r in self.regions:
                s += r.get_sequence()
        else:
            if self.sequence:
                s += self.sequence
            elif self.sequence is None:
                s += "X"
        return s

    def get_len(self) -> Tuple[int, int]:
        """Return the minimum and maximum length of this region."""
        min_l, max_l = 0, 0
        if self.regions:
            for r in self.regions:
                res = r.get_len()
                min_l += res[0]
                max_l += res[1]
        else:
            min_l = self.min_len
            max_l = self.max_len

        return min_l, max_l

    def get_max_length(self):
        """Get the maximum length of this region."""
        return self.get_len()[1]

    def update_attr(self):
        """Update the attributes of the region."""
        if self.regions:
            for idx, r in enumerate(self.regions):
                r.update_attr()

        self.sequence = self.get_sequence()
        self.min_len, self.max_len = self.get_len()
        return

    def __repr__(self) -> str:
        """Return str representation of the instance."""
        d = {
            "region_id": self.region_id,
            "region_type": self.region_type,
            "name": self.name,
            "sequence_type": self.sequence_type,
            "sequence": self.sequence,
            "min_len": self.min_len,
            "max_len": self.max_len,
            "regions": self.regions,
        }
        return f"{d}"

    def to_dict(self):
        """Return a dictionary representation of the region."""
        d = {
            "region_id": self.region_id,
            "region_type": self.region_type,
            "name": self.name,
            "sequence_type": self.sequence_type,
            "sequence": self.sequence,
            "min_len": self.min_len,
            "max_len": self.max_len,
            "regions": [i.to_dict() for i in (self.regions or [])],
        }
        return d

    def get_region_by_id(self, region_id: str) -> Optional["Region"]:
        """Lookup a region by the region_id field."""
        if self.region_id == region_id:
            return self

        if self.regions:
            for r in self.regions:
                if r.region_id == region_id:
                    return r

        return None

    def get_regions_by_type(self, region_type: str) -> List["Region"]:
        """Return all regions with specified region_type."""
        found = []

        if self.region_type == region_type:
            found.append(self)

        if self.regions:
            for r in self.regions:
                found.extend(r.get_regions_by_type(region_type))

        return found

    def get_leaves(self) -> List["Region"]:
        """Return a depth-first list of all leaf regions."""
        leaves = []

        if not self.regions:
            leaves.append(self)
        else:
            for r in self.regions:
                leaves.extend(r.get_leaves())

        return leaves

    def get_leaf_region_types(self) -> Set[RegionType]:
        """Return a set of all leaf region types."""
        leaves = self.get_leaves()
        rtypes = set()

        for r in leaves:
            rtypes.add(r.region_type)

        return rtypes

    def get_subregion_ids(self) -> Set[str]:
        """Return a set of the region_ids of all subregions."""
        ids: Set[str] = set()

        if not self.regions:
            return ids

        for r in self.regions:
            ids.add(r.region_id)
            ids.update(r.get_subregion_ids())

        return ids


class Assay:
    """Class representing a specific assay design.

    An assay contains metadata and a list of (possible nested regions) defining
    the structure of the assay reads.

    :ivar name: Name of the assay
    :ivar assay_spec: List of regions defining the assay structure
    """

    def __init__(self, name: str, assay_spec: Optional[List[Region]] = None):
        """Initialize an assay."""
        self.name = name
        self.assay_spec: List[Region] = assay_spec or []

        self._update_parent_ids()

    @classmethod
    def _validate_region_ids(cls, regions):
        region_ids = set()

        for r in regions:
            if r.region_id in region_ids:
                raise ValueError(f"Duplicate region_id detected: {r.region_id}")

            region_ids.add(r.region_id)

            if r.regions:
                cls._validate_region_ids(r.regions)

    @classmethod
    def _parse_regions(cls, region_data: Optional[List[RegionModel]], level=0):
        # Protected against stack overflow by limiting the number of nested regions
        # 20 nested regions is ridiculously high and should never be reached.
        if level > 20:
            raise RecursionError("Too many nested regions")

        regions: List[Region] = []

        if not region_data:
            return regions

        for r in region_data:
            child_regions = None
            if r.regions:
                child_regions = cls._parse_regions(r.regions, level=level + 1)

            opts = r.dict()
            opts.update({"regions": child_regions})
            region_obj = Region(**opts)

            regions.append(region_obj)

        cls._validate_region_ids(regions)
        return regions

    @classmethod
    def _load_assay_model(cls, assay_data: AssayModel) -> "Assay":
        regions = cls._parse_regions(assay_data.assay_spec)
        assay = Assay(name=assay_data.name, assay_spec=regions)
        return assay

    def _update_parent_ids(self) -> None:
        stack: List[Tuple[Region, Optional[str]]] = [
            (region, None) for region in self.assay_spec
        ]

        while len(stack) > 0:
            region, parent_id = stack.pop()
            if parent_id:
                region.set_parent_id(parent_id)

            if region.regions:
                for r in region.regions:
                    stack.append((r, region.region_id))

    @property
    def region_ids(self) -> Set[str]:
        """Return a set with all regions ids in this assay."""
        ids = set()
        for r in self.assay_spec:
            ids.add(r.region_id)
            ids.update(r.get_subregion_ids())

        return ids

    @classmethod
    def from_yaml(cls, filename: PathType) -> "Assay":
        """Parse an assay from a yaml file.

        :return: an Assay instance loaded from the design config
        :param filename: path to a design config file
        """
        yaml_obj = load_yaml_file(filename)
        checked_obj = AssayModel.model_validate(yaml_obj)
        return cls._load_assay_model(checked_obj)

    def to_json(self):
        """Return a json representation of the assay."""
        return json.dumps(self, default=lambda o: o.__dict__, sort_keys=False, indent=4)

    def get_sequence(self) -> str:
        """Return a sequence representing this region."""
        s = []
        for region in self.assay_spec:
            s.append(region.get_sequence())

        return "".join(s)

    def get_region_by_id(self, region_id: str) -> Optional[Region]:
        """Retrieve a region by its id.

        :param region_id: id of the region to retrieve
        :return: region with the given id or None if not found
        """
        for r in self.assay_spec:
            maybe_r = r.get_region_by_id(region_id)
            if maybe_r:
                return maybe_r

        return None

    @cache
    def get_regions_by_type(self, region_type: RegionType) -> List[Region]:
        """Retrieve all regions of a given type.

        :param region_type: region type to retrieve
        :return: list of regions with the given type
        """
        regions = []
        for r in self.assay_spec:
            regions.extend(r.get_regions_by_type(region_type))

        return regions

    def __len__(self) -> int:
        """Return the total length of the assay in bases."""
        return sum([region.get_max_length() for region in self.assay_spec])


def get_position_in_parent(assay: Assay, region_id: str) -> Tuple[int, int]:
    """Return the start and end position of a region relative to its parent region.

    This assumes the amplicon consists of only fixed length regions in the path
    from the start of the amplicon up until the region of interest.

    :param assay: assay design
    :param region_id: region id of the amplicon
    :return: tuple with start and end position of the region in the parent region
    :raises ValueError: if the region_id is not found in the assay
    :raises ValueError: if the region has no parent
    :raises ValueError: if the parent_region is not found in the assay
    """
    region = assay.get_region_by_id(region_id)
    if region is None:
        raise ValueError(f"Assay does not contain a region with id {region_id}")

    parent_id = region.parent_id
    if parent_id is None:
        raise ValueError(f"Region with id {region_id} has no parent")

    parent_region = assay.get_region_by_id(parent_id)
    if parent_region is None:
        raise ValueError(f"Assay does not contain a region with id {region.parent_id}")

    leaves = parent_region.get_leaves()

    min_start = 0

    for leave in leaves:
        if leave.region_id == region_id:
            break

        min_start += leave.min_len

    min_stop = min_start + region.min_len

    return min_start, min_stop
