"""
This module contains classes and functions related to
the different assay/designs used by pixelator.

Copyright (c) 2022 Pixelgen Technologies AB.
"""
import enum
import json
from typing import Any, List, Mapping, Optional, Set, Tuple
from functools import cache

from pydantic import BaseModel, validator

from pixelator.config.utils import load_yaml_file
from pixelator.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):
    region_id: str
    region_type: RegionType
    name: str
    sequence_type: SequenceType
    sequence: Optional[str]
    min_len: Optional[int]
    max_len: Optional[int]
    data: Optional[Mapping[str, Any]]
    regions: Optional[List["RegionModel"]]

    @validator("sequence")
    def check_valid_dna_string(cls, v: str, values: Mapping[str, Any]) -> str:
        """
        Validate DNA strings.

        :param v: DNA string to validate
        :param values: Pydantic fields
        :raises ValueError: if the string contains non DNA characters
        """
        if values["region_type"] is RegionType.PBS:
            is_valid = all(character in DNA_CHARS for character in v)
            if not is_valid:
                raise ValueError("Not a valid DNA string: found non DNA characters")

        return v


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:
        self.parent_id = 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):
        """
        Set the parent id of this region.

        :param parent_id: parent id to set
        :param recursive: if True, set the parent id of all sub-regions
        """
        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):
        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:
        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):
        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):
        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()
            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.

        :param filename: path to a design config file
        """
        yaml_obj = load_yaml_file(filename)
        checked_obj = AssayModel.parse_obj(yaml_obj)
        return cls._load_assay_model(checked_obj)

    def to_json(self):
        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 update_spec(self):
        for r in self.assay_spec:
            r.update_attr()

    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
    """
    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
