from __future__ import annotations

from enum import StrEnum
from typing import Annotated, ClassVar, Literal

from pydantic import BaseModel, Field

from .common import (
    ArtifactParam,
    ParametersBase,
    ResultBase,
    TableResult,
    TableSourceQuery,
    TableSourceTable,
    TableSourceTaskResult,
)


class OptimizationDirection(StrEnum):
    MAXIMIZE = "MAXIMIZE"
    MINIMIZE = "MINIMIZE"


class ScaleType(StrEnum):
    MULTIPLICATIVE = "MULTIPLICATIVE"
    ADDITIVE = "ADDITIVE"
    RANK = "RANK"


class AllowedMutation(BaseModel):
    in_range: RangeAnnotation = Field(
        description="The annotation expression defining the positions at which the amino acids listed in `amino_acids` can be used."
    )
    amino_acids: list[Annotated[str, Field(description="A single amino acid.")]]


class AnalyzeDataParameters(ParametersBase):
    task_type: ClassVar[str] = "analyze.data/v2"

    dataset: (
        list | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]
    ) = Field(description="The data to analyze before training.")
    reference: MonomerAssembly | VhVlAssembly = Field(
        description="The reference assembly used for computing mutations."
    )
    assays: list[Assay] = Field(description="The assays present in the dataset.")


class AnalyzeDataResult(ResultBase):
    report: ArtifactParam = Field(description="The analysis report.")


class Assay(BaseModel):
    assay_id: str = Field(description="ID of the assay")
    name: str = Field(description="Human readable name of the assay")
    scale_type: ScaleType = Field(
        description="The scale type defines how the assay behaves.\n\n* **Additive**: Assay values are comparable in magnitude over batches, a delta of eg 5 in one batch is equivalent to 5 in another batch.\n* **Multiplicative**: Assay values are comparable over batches by applying a multiplication factor, eg fold improvement: `score_variantB_batch1 = factor(typically starting sequence score) * score_variantB_batch`.\n* **Rank**: Assay values are not comparable over batches, we can only assume the ranking within a batch is correct."
    )
    unit: str | None = Field(default=None, description="Unit of the assay - used for display purposes only")


class AssemblyFeature(BaseModel):
    ranges: dict[str, list[tuple[int, int]]] = Field(description="The ranges of the feature, keyed by polymer name")
    metadata: dict[str, str | int | float | bool] | None = Field(
        default=None, description="Optional metadata for the feature"
    )


class CombinatorialGenerator(BaseModel):
    type: Literal["COMBINATORIAL"] = Field(default="COMBINATORIAL")
    template: MonomerAssembly | VhVlAssembly = Field(
        description="The assembly that serves as template for generation, including any annotations to be used in `requirements`."
    )
    max_mutations: int = Field(
        description="Maximum number of mutations to combine in generation. Lower values will result in fewer total candidate sequences evaluated for this generator."
    )
    requirements: list[Requirement] = Field(
        default_factory=list,
        description="A list of requirements to satisfy. All requirements must be met for a valid assembly. See the `Requirement` documentation for details on how requirements are evaluated.",
    )
    allowed_mutations: list[AllowedMutation] = Field(
        description="The complete set of mutations that can be used in generation."
    )
    alias: str = Field(
        default="combinatorial_generator",
        description="Custom alias to refer to this generator within the task. Useful for tracking provenance of assemblies generated by this generator. For example, this name will appear in the `generator_provenance` column for rows in `selected_sequences` which were produced by this generator.",
    )


class Constraint(BaseModel):
    measure: Measure = Field(description="The measure to constrain.")
    threshold: float | RelativeTo = Field(
        description="The threshold for the constraint, either an absolute value or relative to a reference."
    )


class DiversifyParameters(ParametersBase):
    task_type: ClassVar[str] = "diversify/v2"

    generator: Sequential | Parallel | ModelBasedGenerator | Explicit | CombinatorialGenerator = Field(
        description="The generator configuration or configurations that define requirements and strategies for generating candidate assemblies.",
        discriminator="type",
    )
    ranker: LikelihoodBasedRanker = Field(
        description="The configuration defining how candidate assemblies are evaluated for inclusion in `selected_assemblies`."
    )
    num_assemblies: int = Field(
        description="The target number of `selected_assemblies` to output. In rare cases, the actual number of output assemblies may be slightly lower."
    )


class DiversifyResult(ResultBase):
    selected_assemblies: TableResult = Field(
        description="Table containing the selected assemblies output by `diversify`."
    )
    report: ArtifactParam | None = Field(description="A report on the result of the task.")


class EngineerParameters(ParametersBase):
    task_type: ClassVar[str] = "engineer/v2"

    generator: Sequential | Parallel | ModelBasedGenerator | Explicit | CombinatorialGenerator = Field(
        description="The generator configuration or configurations that define requirements and strategies for generating candidate assemblies.",
        discriminator="type",
    )
    ranker: ScoreBasedRanker = Field(
        description="The configuration defining how candidate assemblies are evaluated for inclusion in `selected_assemblies`."
    )
    num_assemblies: int = Field(
        description="The target number of selected assemblies to output. In rare cases of very difficult `constraints` this task may produce a smaller number."
    )


class EngineerResult(ResultBase):
    selected_assemblies: TableResult = Field(
        description="Table containing the selected assemblies output by `engineer`."
    )
    report: ArtifactParam | None = Field(description="A report on the result of the task.")


class Explicit(BaseModel):
    type: Literal["EXPLICIT"] = Field(default="EXPLICIT")
    source: MonomerSource | VhVlSource
    alias: str = Field(
        default="explicit_generator",
        description="Custom alias to refer to this generator within the task. Useful for tracking provenance of assemblies generated by this generator. For example, this name will appear in the `generator_provenance` column for rows in `selected_sequences` which were produced by this generator.",
    )


class LikelihoodBasedRanker(BaseModel):
    type: Literal["LIKELIHOOD_BASED"] = Field(default="LIKELIHOOD_BASED")
    sampler_config: MonomerSamplerConfig | VhVlSamplerConfig = Field(
        description="The trained sampler artifact used to compute assembly likelihoods, along with any optional control assemblies to include in the final output for experimental validation.",
        discriminator="kind",
    )
    hit_redundancy: int | None = Field(
        default=None,
        description='The maximum number of similar variants within the proposed set.\n\nWe say two variants are "similar" when their predicted performances are highly correlated.\nAt lower values this will mean on average fewer hits, but greater diversity within the set,\nwhile at maximum value (`hit_redundancy=num_selected`) the selection is purely greedy, with\nno requirement on diversity within the set. If set to None, it will be automatically set\nto a default value of `1 + num_assemblies//20`.',
    )


class Measure(BaseModel):
    assay_id: str = Field(description="The ID of the assay.")
    direction: OptimizationDirection = Field(description="The direction of optimization for the assay.")


class ModelBasedGenerator(BaseModel):
    type: Literal["MODEL_BASED"] = Field(default="MODEL_BASED")
    template: MonomerSamplingTemplate | VhVlSamplingTemplate = Field(
        description="The template sequence to use for generation.", discriminator="kind"
    )
    requirements: list[Requirement] = Field(
        default_factory=list,
        description="A list of requirements to satisfy. All requirements must be met for a valid assembly. See the `Requirement` documentation for details on how requirements are evaluated.",
    )
    discourage_mutations: RangeAnnotation | None = Field(
        default=None, description="Expression defining regions where mutations should be discouraged."
    )
    min_mutations: int = Field(description="Minimum number of mutations from the template assembly.")
    max_mutations: int = Field(description="Maximum number of mutations from the template assembly.")
    alias: str = Field(
        default="model_based_generator",
        description="Custom alias to refer to this generator within the task. Useful for tracking provenance of assemblies generated by this generator. For example, this name will appear in the `generator_provenance` column for rows in `selected_sequences` which were produced by this generator.",
    )


class MonomerAssembly(BaseModel):
    monomer: str = Field(description="The monomer sequence")


class MonomerData(BaseModel):
    kind: Literal["MONOMER"] = Field(default="MONOMER")
    homologs: (
        list | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]
    ) = Field(description="The homologous assemblies used for training.")
    dataset: (
        list
        | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]
        | None
    ) = Field(
        description="The assayed assemblies used for training. If None, predictor finetuning and generator conditioning will be skipped."
    )


class MonomerModels(BaseModel):
    kind: Literal["MONOMER"] = Field(default="MONOMER")
    sampler: ArtifactParam = Field(description="The monomer sampler model.")
    scorer: ArtifactParam | None = Field(description="The monomer scorer model.")
    conditioned_sampler: ArtifactParam | None = Field(
        description="The monomer sampler model conditioned on the primary objective."
    )


class MonomerSamplerConfig(BaseModel):
    kind: Literal["MONOMER_SAMPLER"] = Field(default="MONOMER_SAMPLER")
    sampler: ArtifactParam = Field(description="The trained sampler artifact used to compute sequence likelihoods.")
    reference_assembly: MonomerAssembly = Field(
        description="The reference assembly used as the denominator when computing sequence likelihood ratios."
    )
    controls: (
        list
        | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]
        | None
    ) = Field(
        default=None,
        description="Optional control assemblies to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_assemblies` will be intentionally diverse from these assemblies. However these assemblies will *not* be included in the `selected_assemblies` output table, so -- if used -- take care to re-include them. ",
    )


class MonomerSamplingTemplate(BaseModel):
    kind: Literal["MONOMER"] = Field(default="MONOMER")
    assembly: MonomerAssembly = Field(description="The template sequence to use for generation.")
    sampler: ArtifactParam = Field(description="The trained sampler artifact used to generate candidate sequences.")


class MonomerScorerConfig(BaseModel):
    kind: Literal["MONOMER_SCORER"] = Field(default="MONOMER_SCORER")
    scorer: ArtifactParam = Field(description="The trained scoring model artifact used to score sequences.")
    controls: (
        list
        | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]
        | None
    ) = Field(
        default=None,
        description="Optional control sequences to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_sequences` will be intentionally diverse from these sequences. However these sequences will *not* be included in the `selected_sequences` output table, so -- if used -- take care to re-include them. ",
    )


class MonomerSource(BaseModel):
    kind: Literal["MONOMER"] = Field(default="MONOMER")
    source: list | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]


class OASDatabaseSearch(BaseModel):
    type_: Literal["OAS_SEARCH"] = Field(default="OAS_SEARCH", alias="type")
    seed_assemblies: list[VhVlAssembly] = Field(description="The list of light/heavy chain assemblies to search for.")
    search_unpaired: bool = Field(default=True, description="Whether to search for unpaired heavy/light chains.")
    search_paired: bool = Field(default=True, description="Whether to search for paired heavy/light chains.")
    annotate_assembly_features: bool = Field(
        default=True, description="Whether to infer CDR features based on OAS annotations."
    )


class Parallel(BaseModel):
    type: Literal["PARALLEL"] = Field(default="PARALLEL")
    ensemble: list[
        Annotated[
            Sequential | Parallel | ModelBasedGenerator | Explicit | CombinatorialGenerator,
            Field(description="The generation strategy for creating candidate sequences.", discriminator="type"),
        ]
    ] = Field(default_factory=list, description="The list of generators to run in parallel within this ensemble.")
    alias: str | None = Field(
        default=None,
        description="Optional alias to refer to this generator within the task. Useful for tracking provenance of assemblies generated by this generator. If provided, this name will appear in the `generator_provenance` column for rows in `selected_sequences` which were produced by this generator.",
    )


class PreRanked(BaseModel):
    type: Literal["PRE_RANKED"] = Field(default="PRE_RANKED")
    hit_redundancy: int = Field(
        default=3,
        description='The maximum number of similar variants within the proposed set.\n\nWe say two variants are "similar" when their predicted performances are highly correlated.\nAt lower values this will mean on average fewer hits, but greater diversity within the set,\nwhile at maximum value (`hit_redundancy=num_assemblies`) the selection is purely greedy, with\nno requirement on diversity within the set.',
    )


class PrimaryObjective(BaseModel):
    measure: Measure = Field(description="The measure to optimize.")
    reference: float | RelativeTo = Field(
        description="The reference relative to which the primary objective is defined. For multiplicative assays, an absolute reference has to be positive."
    )


class RangeAnnotation(BaseModel):
    ranges: dict[str, list[tuple[int, int]]] = Field(
        description="A mapping from polymer names to annotation ranges. Each range is a tuple of (start, end). Indices are 0-based and inclusive for the start and exclusive for the end. For example, (0, 5) refers to positions 0 to 4."
    )


class RelativeTo(BaseModel):
    reference: MonomerAssembly | VhVlAssembly = Field(description="The reference assembly.")
    margin: float | None = Field(
        default=None,
        description="The desired improvement over the `reference` assembly.\n\nRegardless of the `OptimizationDirection` of the `Measure` that this `RelativeTo` is referring to, larger numbers always mean higher desired improvements.\n\nFor additive assays, a positive margin corresponds to an improvement over the `reference`, and a negative margin to a reduction in performance. For multiplicative assays, the margin is the expected fold improvement, with values greater than 1.0 corresponding to an improvement over the `reference`, and values between 0.0 and 1.0 corresponding to a reduction in performance; negative values are not allowed.\n\nMathematically, the constraint is computed as follows for different scale types and optimization directions:\n\n| Scale Type       | MAXIMIZE | MINIMIZE |\n| ---------------- | -------------------------------- | -------------------------------- |\n| Additive         | `reference + margin`             | `reference - margin`             |\n| Multiplicative   | `reference * margin`             | `reference / margin`             |\n\nFor example, a 10% improvement for a multiplicative assay corresponds to a margin of 1.1.",
    )


class Requirement(BaseModel):
    apply_to: RangeAnnotation
    motifs: list[Annotated[str, Field(description="A motif or regular expression of amino acids.")]]
    operator: Literal["none", "all", "any"] = Field(
        description="The choice of how to interpret the `motifs` together as a requirement."
    )


class ScoreBasedRanker(BaseModel):
    type: Literal["SCORE_BASED"] = Field(default="SCORE_BASED")
    primary_objective: PrimaryObjective = Field(description="The primary optimization objective of `engineer`.")
    constraints: list[Constraint] = Field(
        default_factory=list, description="List of constraints `engineer` will attempt to satisfy."
    )
    scorer_config: MonomerScorerConfig | VhVlScorerConfig = Field(
        description="The trained scoring model artifact used to score sequences, along with any optional control sequences to include in the final output for experimental validation.",
        discriminator="kind",
    )
    hit_redundancy: int | None = Field(
        default=None,
        description='The maximum number of similar variants within the proposed set.\n\nWe say two variants are "similar" when their predicted performances are highly correlated.\nAt lower values this will mean on average fewer hits, but greater diversity within the set,\nwhile at maximum value (`hit_redundancy=num_selected`) the selection is purely greedy, with\nno requirement on diversity within the set. If set to None, it will be automatically set\nto a default value of `1 + num_assemblies//20`.',
    )


class SearchParameters(ParametersBase):
    task_type: ClassVar[str] = "search/v2"

    source: UnirefDatabaseSearch | OASDatabaseSearch = Field(
        description="The source database to search against, along with the assemblies to search for, and database-specific configuration.",
        discriminator="type_",
    )


class SearchResult(ResultBase):
    assembly_features: list[dict[str, AssemblyFeature]] = Field(
        default_factory=list,
        description="Parallel array (wrt to `seed_assemblies`) of features, keyed by feature name.",
    )
    homologs: TableResult
    report: ArtifactParam | None = Field(default=None, description="A report on the result of the task.")


class SelectParameters(ParametersBase):
    task_type: ClassVar[str] = "select/v2"

    assemblies: MonomerSource | VhVlSource = Field(description="The source of assemblies to select from.")
    ranker: (
        Annotated[
            LikelihoodBasedRanker | ScoreBasedRanker,
            Field(description="The ranking strategy for selecting candidate assemblies.", discriminator="type"),
        ]
        | PreRanked
    ) = Field(
        description="The configuration defining how candidate assemblies are evaluated for inclusion in `selected_assemblies`."
    )
    min_assemblies_to_select: int = Field(
        description="The minimum number of assemblies to select. The number of selected candidates is dynamically determined and may be greater than this number. If an exact number of assemblies is desired, set `min_assemblies_to_select` and `max_assemblies_to_select` to the same value."
    )
    max_assemblies_to_select: int | None = Field(
        default=None,
        description="An optional (inclusive) upper limit on the number of assemblies to select. If provided, the task will select between `min_assemblies_to_select` and `max_assemblies_to_select` assemblies.",
    )


class SelectResult(ResultBase):
    selected_assemblies: TableResult = Field(
        description="Table containing the selected assemblies output by `select` sorted in descending order of preference."
    )


class Sequential(BaseModel):
    type: Literal["SEQUENTIAL"] = Field(default="SEQUENTIAL")
    ensemble: list[
        Annotated[
            Sequential | Parallel | ModelBasedGenerator | Explicit | CombinatorialGenerator,
            Field(description="The generation strategy for creating candidate sequences.", discriminator="type"),
        ]
    ] = Field(default_factory=list, description="The list of generators to run in sequence within this ensemble.")
    alias: str | None = Field(
        default=None,
        description="Optional alias to refer to this generator within the task. Useful for tracking provenance of assemblies generated by this generator. If provided, this name will appear in the `generator_provenance` column for rows in `selected_sequences` which were produced by this generator.",
    )


class TrainParameters(ParametersBase):
    task_type: ClassVar[str] = "train/v2"

    data: MonomerData | VhVlData = Field(
        description="The training data, including homologs and assayed sequences.", discriminator="kind"
    )
    assays: list[Assay] = Field(
        default_factory=list,
        description="The assays used for sampler conditioning (on the `primary_objective`) and scorer finetuning and evaluation. Must correspond to assays present in `data.dataset`. If `data.dataset` is None, `assays` must be empty. Otherwise, it must contain at least one assay.",
    )
    primary_objective: Measure | None = Field(
        default=None,
        description="The assay on which the sampler model is conditioned. Should correspond to an assay in `dataset`. If `data.dataset` is None, `primary_objective` must not be set. Otherwise, it must be set.",
    )


class TrainResult(ResultBase):
    models: MonomerModels | VhVlModels = Field(description="The trained models.", discriminator="kind")
    report: ArtifactParam | None = Field(description="A report on the result of the task.")


class UnirefDatabaseSearch(BaseModel):
    type_: Literal["UNIREF_SEARCH"] = Field(default="UNIREF_SEARCH", alias="type")
    seed_assemblies: list[MonomerAssembly] = Field(description="A list of monomeric assemblies to search for.")


class VhVlAssembly(BaseModel):
    vh: str | None = Field(default=None, description="The heavy chain sequence.")
    vl: str | None = Field(default=None, description="The light chain sequence.")


class VhVlData(BaseModel):
    kind: Literal["VHVL"] = Field(default="VHVL")
    homologs: (
        list | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]
    ) = Field(description="The homologous assemblies used for training.")
    dataset: (
        list
        | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]
        | None
    ) = Field(
        description="The assayed assemblies used for training. If None, predictor finetuning and generator conditioning will be skipped."
    )


class VhVlModels(BaseModel):
    kind: Literal["VH_VL"] = Field(default="VH_VL")
    sampler: ArtifactParam = Field(description="The VhVl sampler model.")
    scorer: ArtifactParam | None = Field(description="The VhVl scorer model.")
    conditioned_sampler: ArtifactParam | None = Field(
        description="The VhVl sampler model conditioned on the primary objective."
    )


class VhVlSamplerConfig(BaseModel):
    kind: Literal["VHVL_SAMPLER"] = Field(default="VHVL_SAMPLER")
    sampler: ArtifactParam = Field(description="The trained sampler artifact used to compute assembly likelihoods.")
    reference_assembly: VhVlAssembly = Field(
        description="The reference assembly used as the denominator when computing assembly likelihood ratios."
    )
    controls: (
        list
        | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]
        | None
    ) = Field(
        default=None,
        description="Optional control assemblies to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_assemblies` will be intentionally diverse from these assemblies. However these assemblies will *not* be included in the `selected_assemblies` output table, so -- if used -- take care to re-include them. ",
    )


class VhVlSamplingTemplate(BaseModel):
    kind: Literal["VHVL"] = Field(default="VHVL")
    assembly: VhVlAssembly = Field(description="The template sequence to use for generation.")
    sampler: ArtifactParam = Field(description="The trained sampler artifact used to generate candidate sequences.")


class VhVlScorerConfig(BaseModel):
    kind: Literal["VHVL_SCORER"] = Field(default="VHVL_SCORER")
    scorer: ArtifactParam = Field(description="The trained scoring model artifact used to score sequences.")
    controls: (
        list
        | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]
        | None
    ) = Field(
        default=None,
        description="Optional control sequences to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_sequences` will be intentionally diverse from these sequences. However these sequences will *not* be included in the `selected_sequences` output table, so -- if used -- take care to re-include them. ",
    )


class VhVlSource(BaseModel):
    kind: Literal["VHVL"] = Field(default="VHVL")
    source: list | Annotated[TableSourceTable | TableSourceQuery | TableSourceTaskResult, Field(discriminator="kind")]
