"""The types used by our R/qtl2 system."""
from enum import Enum
from pathlib import Path
from dataclasses import field, dataclass
from typing import (Union,
                    Literal,
                    Optional,
                    Sequence,
                    TypeAlias,
                    Annotated)


_ALLELE_COUNTS_BY_CROSSTYPE = {
    "bc": 2, "f2": 2, "riself": 2, "risib": 2, "dh": 2, "haploid": 2, "ail": 2,
    "ail3": 3,
    "riself4": 4, "risib4": 4,
    "dh6": 6,
    "hs": 8, "do": 8, "riself8": 8, "risib8": 8,
    "dof1": 9, # 8 DO founders + 1 base inbred line
    "riself16": 16,
    "magic19": 19
}



PathLike = Union[Path, str]

AlleleLabel: TypeAlias = str

GenoCodeAlleleValue = Union[
    Literal[-1, 0, 1],
    Literal[1, 2, 3],# T0D0: reconcile this to genocodes of 0, 1, 2 instead
    Annotated[float, "Value must be between 0.0 and 2.0"]
]

GenoCode = dict[AlleleLabel, GenoCodeAlleleValue]

# Sex information
SexKey = Union[str, int]# codes used in the covariate file
@dataclass(frozen=True)
class ControlDataSex:
    """Sex information in the control file."""
    covar: Optional[str] = None
    file: Optional[PathLike] = None
    female: Optional[SexKey] = None
    male: Optional[SexKey] = None


# Cross information
CrossType = Literal[
    # T0D0: Look at https://kbroman.org/qtl2/assets/vignettes/input_files.html#Backcross to implement appropriate QC dependent on the crosstype.
    "bc",# Backcross
    "f2",# F2 intercross
    "riself",# RIL by self -- RIL: Recombinant Inbred Lines
    "risib",# RIL by sib
    "dh",# Double haploid
    "haploid",# Haploid
    "ail",# AIL -- Advanced Inbred Lines
    "hs",# Heterogeneous stock
    "do",# Diversity outbreds
    "riself4",# Multi-parent recombinant inbred lines
    "riself8",# Multi-parent recombinant inbred lines
    "riself16",# Multi-parent recombinant inbred lines
    "risib4",# Multi-parent recombinant inbred lines
    "risib8",# Multi-parent recombinant inbred lines
    "magic19",# 19-way MAGIC lines
    "dh6",# 6-way doubled haploids
    "dof1",# DOF1
    "ail3",# 3-way advanced intercross lines
    "genail",# General advanced intercross lines
    "genril",# General recombinant inbred lines
]

class CrossDirection(Enum):
    FORWARD = 0
    REVERSE = 1


@dataclass(frozen=True)
class ShortCrossInfo:
    """Cross information in the control file."""
    # "covar": indicates the name of the column in the covariate data
    covar: Optional[str] = None
    file: Optional[PathLike] = None
    forward_crosses: Sequence[str] = field(default_factory=tuple)
    reverse_crosses: Sequence[str] = field(default_factory=tuple)

    def __post_init__(self):
        """Validate the cross info."""
        # A valid short config MUST have either a `covar` or a `file` source
        # but not both.
        if not (self.covar and self.file):
            raise ValueError("ShortCrossInfo requires either a 'covar' or a "
                             "'file' parameter.")
        if self.covar and self.file:
            raise ValueError("ShortCrossInfo cannot define both 'covar' and "
                             "'file' simultaneously.")


LongCrossInfo: TypeAlias = PathLike
CrossInfo = Union[LongCrossInfo, ShortCrossInfo]


def expected_allele_count_for_crosstype(
        alleles: Sequence[str],
        crosstype: CrossType,
        expected_counts: dict[CrossType, int]
) -> bool:
    """Check that the number of alleles matches what crosstype expects."""
    def __validate__(expected, alleles):
        if expected is not None and len(alleles) != expected:
            raise ValueError(
                f"Cross type '{crosstype}' expects exactly {expected} allele "
                f"labels. Received {len(alleles)}: {alleles}")

    if crosstype.startswith(("genail", "genril")):
        # genail and genril are dynamic (e.g., genail8)
        # Extract the trailing digits to determine expected founder count
        num_part = "".join(filter(str.isdigit, crosstype))
        if num_part:
            return __validate__(int(num_part), alleles)

    return __validate__(expected_counts.get(crosstype), alleles)
            


@dataclass(frozen=True)
class ControlData:
    """Class for the R/qtl2 control data."""
    # File names: Force listings, rather than singular strings
    geno: Sequence[PathLike] = field(default_factory=tuple)
    founder_geno: Sequence[PathLike] = field(default_factory=tuple)
    pheno: Sequence[PathLike] = field(default_factory=tuple)
    covar: Sequence[PathLike] = field(default_factory=tuple)
    phenocovar: Sequence[PathLike] = field(default_factory=tuple)
    gmap: Sequence[PathLike] = field(default_factory=tuple)
    pmap: Sequence[PathLike] = field(default_factory=tuple)

    # X Chromosome
    x_chr: Optional[str] = None

    # Allele labels
    alleles: Sequence[str] = field(default_factory=tuple)

    # Genotype codes
    genotypes: GenoCode = field(default_factory=dict)

    # sex
    sex: Optional[ControlDataSex] = None

    # Cross info
    crosstype: Optional[CrossType] = None
    cross_info: Optional[CrossInfo] = None


    # CSV fields
    na_strings: Sequence[str] = ("-", "NA", "N/A")
    sep: str = ","
    comment_char: Optional[str] = "#"

    def __post_init__(self):
        if self.alleles and self.crosstype:
            expected_allele_count_for_crosstype(
                self.alleles, self.crosstype, _ALLELE_COUNTS_BY_CROSSTYPE)