from __future__ import annotations

from pathlib import Path
from typing import TypedDict

import pandas as pd  # type: ignore
from sqlalchemy import Boolean
from sqlalchemy import create_engine
from sqlalchemy import Float
from sqlalchemy import ForeignKey
from sqlalchemy import func
from sqlalchemy import inspect
from sqlalchemy import Integer
from sqlalchemy import JSON
from sqlalchemy import select
from sqlalchemy import String
from sqlalchemy import Text
from sqlalchemy.engine import Connectable
from sqlalchemy.engine import Engine
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column
from sqlalchemy.orm import object_session
from sqlalchemy.orm import relationship
from sqlalchemy.orm import with_parent

from .utils import check_missing_columns
from .utils import secondary

# func.count()
# pylint: disable=not-callable


class Base(DeclarativeBase):
    pass


def clear_orm(Orm: type) -> None:
    if hasattr(Orm, "__table__"):
        Base.metadata.remove(getattr(Orm, "__table__"))
    elif hasattr(Orm, "__tablename__"):
        tname = getattr(Orm, "__tablename__")
        if tname in Base.metadata.tables:
            Base.metadata.remove(Base.metadata.tables[tname])
    Base.registry.dispose()


class Modification(TypedDict):
    mass: float
    massdiff: float | None
    position: int
    site: str | None


class Peptide(Base):
    __tablename__ = "peptides"
    peptideid: Mapped[int] = mapped_column(Integer, primary_key=True)
    peptide: Mapped[str] = mapped_column(Text, nullable=False)
    assumed_charge: Mapped[int | None] = mapped_column(Integer)
    modcol: Mapped[str | None] = mapped_column(Text)
    calc_neutral_pep_mass: Mapped[float | None] = mapped_column(Float)
    precursor_neutral_mass: Mapped[float | None] = mapped_column(Float)
    peptideprophet_probability: Mapped[float | None] = mapped_column(Float)
    mz: Mapped[float | None] = mapped_column(Float)
    observed_mz: Mapped[float | None] = mapped_column(Float)
    retention_time_sec: Mapped[float | None] = mapped_column(Float)
    num_missed_cleavages: Mapped[int | None] = mapped_column(Integer)
    is_decoy: Mapped[bool | None] = mapped_column(Boolean)
    agg_count: Mapped[int | None] = mapped_column(Integer)
    # proteins: Mapped[list[str]] = mapped_column(JSON)
    protein_descr: Mapped[list[str] | None] = mapped_column(JSON)
    modifications: Mapped[list[Modification] | None] = mapped_column(JSON)
    run: Mapped[str | None] = mapped_column(String(64))
    isLabelledRun: Mapped[bool | None] = mapped_column(Boolean)
    fdr: Mapped[float | None] = mapped_column(Float)
    # probability = Column(Float)
    # percent_coverage = Column(Float)
    # confidence = Column(Float)
    # ngroups = Column(Integer)
    # prStr = Column(Text)
    # description = Column(Text)
    # prGrassumed_chargeoup = Column(Integer)
    # ----- computed -----
    labelledElementCount: Mapped[int | None] = mapped_column(Integer)
    maxIso: Mapped[int | None] = mapped_column(Integer)
    mzranges: Mapped[list[float] | None] = mapped_column(JSON, deferred=True)
    eics: Mapped[list[float] | None] = mapped_column(JSON, deferred=True)
    # isotopeRegressions = Column(JSON)
    adjustedRsq: Mapped[list[float] | None] = mapped_column(
        JSON,
        deferred=True,
    )  # list[float]
    isotopeEnvelopes: Mapped[list[float] | None] = mapped_column(JSON, deferred=True)
    monoPeakArea: Mapped[float | None] = mapped_column(Float)
    maxPeakArea: Mapped[float | None] = mapped_column(Float)
    monoFitParams: Mapped[list[float] | None] = mapped_column(JSON, deferred=True)
    inv_ratio: Mapped[float | None] = mapped_column(Float)
    # heavyDistribution = Column(JSON)
    # heavyFitParams = Column(JSON)
    # heavy_adj_r_squared = Column(Float)
    labelledEnvelopes: Mapped[list[float] | None] = mapped_column(JSON, deferred=True)
    relativeIsotopeAbundance: Mapped[float | None] = mapped_column(Float)
    enrichment: Mapped[float | None] = mapped_column(Float)
    labelEnrichment2: Mapped[float | None] = mapped_column(Float)
    heavyCor: Mapped[float | None] = mapped_column(Float)
    heavyCor2: Mapped[float | None] = mapped_column(Float)
    theoreticalDist: Mapped[list[float] | None] = mapped_column(JSON, deferred=True)
    nnls_residual: Mapped[float | None] = mapped_column(Float)
    totalNNLSWeight: Mapped[float] = mapped_column(Float)
    fail: Mapped[int] = mapped_column(Integer)
    eics_shape: Mapped[list[int] | None] = mapped_column(JSON, deferred=True)
    protein_names: Mapped[list[str] | None] = mapped_column(JSON)

    proteins: Mapped[list[Protein]] = relationship(
        lambda: Protein,
        secondary="peptides_proteins",
        back_populates="peptides",
    )

    @hybrid_property
    def norm_nnls_residual(self) -> float:
        if self.totalNNLSWeight == 0.0:
            return 0.0
        if self.nnls_residual is None:
            return 0.0
        return self.nnls_residual / self.totalNNLSWeight

    @norm_nnls_residual.expression
    def _norm_nnls_residual(cls):  # pylint: disable=no-self-argument
        return cls.nnls_residual / cls.totalNNLSWeight

    @hybrid_property
    def protein_count(self):
        q = select(func.count()).where(with_parent(self, Peptide.proteins))
        return object_session(self).execute(q).scalar_one()

    @protein_count.expression
    def _protein_count_expression(cls):  # pylint: disable=no-self-argument
        q = (
            select(func.count(PepProt.peptideid))
            .where(PepProt.peptideid == cls.peptideid)
            .label("protein_count")
        )
        return q

    def __repr__(self):
        return f"Peptide({self.peptide},{self.mz}, {self.assumed_charge})"

    iso_peaks_nr_20: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        server_default="0",
    )
    iso_peaks_nr_40: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        server_default="0",
    )
    iso_peaks_nr_60: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        server_default="0",
    )
    iso_peaks_nr_80: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        server_default="0",
    )


MINRSQ100 = [20, 40, 60, 80]


class PI(TypedDict):
    protein_name: str


class PeptideInfo(TypedDict):
    calc_neutral_pep_mass: float
    charge: int
    exp_tot_instances: float
    fpkm_adjusted_probability: float
    group_weight: str  # str of float
    initial_probability: float
    is_contributing_evidence: bool
    is_nondegenerate_evidence: bool
    mod_aminoacid_mass: float | None
    mod_nterm_mass: float | None
    modified_peptide: str | None
    n_enzymatic_termini: int
    n_instances: int
    n_sibling_peptides: float
    n_sibling_peptides_bin: int
    nsp_adjusted_probability: float
    peptide_group_designator: str
    peptide_parent_protein: list[PI] | None
    peptide_sequence: str
    weight: float


class Protein(Base):
    __tablename__ = "proteins"

    proteinid: Mapped[int] = mapped_column(Integer, primary_key=True)
    probability: Mapped[int | None] = mapped_column(Float)
    group_number: Mapped[int | None] = mapped_column(Integer)
    protein_name: Mapped[str | None] = mapped_column(Text)
    n_indistinguishable_proteins: Mapped[int | None] = mapped_column(Integer)
    percent_coverage: Mapped[float | None] = mapped_column(Float)
    unique_stripped_peptides: Mapped[list[str] | None] = mapped_column(JSON)
    group_sibling_id: Mapped[str | None] = mapped_column(String(16))
    total_number_peptides: Mapped[int | None] = mapped_column(Integer)
    total_number_distinct_peptides: Mapped[int | None] = mapped_column(Integer)
    confidence: Mapped[float | None] = mapped_column(Float)
    peptide: Mapped[list[PeptideInfo] | None] = mapped_column(JSON)
    protein_description: Mapped[str | None] = mapped_column(Text)
    pct_spectrum_ids: Mapped[float | None] = mapped_column(Float)
    prot_length: Mapped[int | None] = mapped_column(Integer)
    indistinguishable_protein: Mapped[list[str] | None] = mapped_column(JSON)
    lpf_median: Mapped[float | None] = mapped_column(Float)
    lpf_std: Mapped[float | None] = mapped_column(Float)
    num_found_peptides: Mapped[int | None] = mapped_column(Integer)
    # subsuming_protein_entry = Column(JSON)
    peptides: Mapped[list[Peptide]] = relationship(
        Peptide,
        secondary="peptides_proteins",
        back_populates="proteins",
    )

    @hybrid_property
    def peptide_count(self) -> int:
        sess = object_session(self)
        if sess is None:
            raise RuntimeError("no session for Peptide")
        q = select(func.count()).where(with_parent(self, Protein.peptides))
        return sess.execute(q).scalar_one()

    @peptide_count.expression
    def _peptide_count_expression(cls):  # pylint: disable=no-self-argument
        q = (
            select(func.count())
            .where(PepProt.proteinid == cls.proteinid)
            .label("peptide_count")
        )
        return q

    def __repr__(self) -> str:
        return f"Protein({self.protein_name},prob={self.probability},npep={self.unique_stripped_peptides})"


class PepProt(Base):
    __tablename__ = "peptides_proteins"
    peptideid: Mapped[int] = mapped_column(
        Integer,
        ForeignKey(Peptide.peptideid),
        nullable=False,
        primary_key=True,
    )
    proteinid: Mapped[int] = mapped_column(
        Integer,
        ForeignKey(Protein.proteinid),
        nullable=False,
        primary_key=True,
    )


def aggregate_lpf(
    prot: pd.DataFrame,
    pep: pd.DataFrame,
    secondary: pd.DataFrame,
) -> pd.DataFrame:
    m = secondary.merge(pep, how="inner", left_on="peptideid", right_on="peptideid")
    a = (
        m.groupby("proteinid")
        .agg(
            lpf_median=pd.NamedAgg(column="relativeIsotopeAbundance", aggfunc="median"),
            lpf_std=pd.NamedAgg(column="relativeIsotopeAbundance", aggfunc="std"),
            num_found_peptides=pd.NamedAgg(column="peptideid", aggfunc="count"),
        )
        .reset_index()
    )
    a.fillna(0.0, inplace=True)

    return prot.merge(a, how="left", left_on="proteinid", right_on="proteinid")


def save_db(
    pep: pd.DataFrame,
    prot: pd.DataFrame,
    engine: Engine,
) -> tuple[int, int, int]:
    from .utils import dehydrate_peptides, dehydrate_prot
    from .iso_peaks import add_iso_peaks

    # from .precompute import populate_isoPeaksNrFixed
    Base.metadata.create_all(bind=engine)

    pep = dehydrate_peptides(pep)
    prot = dehydrate_prot(prot)
    # pep["peptideid"] = pep.index + 1
    prot["proteinid"] = prot.index + 1
    sec, names = secondary(pep, prot)
    # add protein_names concatenation to peptide
    pep = pep.merge(names, how="left", left_on="peptideid", right_on="peptideid")
    # add iso_peaks_nrNN peaks
    pep = add_iso_peaks(pep)
    # don't save any proteins that don't exists
    prot = prot[prot["proteinid"].isin(sec["proteinid"])]

    prot = aggregate_lpf(prot, pep, sec)
    with engine.begin() as conn:
        npep = save_df(pep, Peptide, conn)
        nprot = save_df(prot, Protein, conn)
        nsec = save_df(sec, PepProt, conn)

    # populate_isoPeaksNrFixed(engine, force=True)
    return npep, nprot, nsec


def save_df(
    df: pd.DataFrame,
    Orm: type[Base],
    conn: Connectable,
    chunksize: int = 1000,
) -> int:
    i = inspect(Orm)
    dtypes = {c.name: c.type for c in i.columns}

    return df[list(dtypes)].to_sql(
        i.local_table.name,  # type: ignore
        con=conn,
        # need this for some reason....
        dtype=dtypes,
        chunksize=chunksize,
        if_exists="append",
        index=False,
    )


def check_missing_pep(pep: pd.DataFrame) -> set[str]:
    return check_missing_columns(pep, Peptide)


def check_missing_prot(prot: pd.DataFrame) -> set[str]:
    return check_missing_columns(prot, Protein)


def file2engine(filename: str | Path, *, echo: bool = False) -> Engine:
    return create_engine(f"sqlite:///{filename}", future=True, echo=echo)


def save_to_file(
    sqlfile: Path,
    pep: pd.DataFrame,
    prot: pd.DataFrame,
) -> tuple[int, int, int]:
    if sqlfile.exists():
        sqlfile.unlink()

    engine = file2engine(sqlfile)
    try:
        return save_db(pep, prot, engine)
    finally:
        engine.dispose()


if __name__ == "__main__":
    import click

    @click.command()
    @click.argument("database")
    def dbtest(database: str) -> None:
        db = Path(database)
        if db.is_file():
            db.unlink(missing_ok=True)
        click.secho(f"creating {database}")
        engine = file2engine(database)
        Base.metadata.create_all(bind=engine)

    dbtest()  # pylint: disable=no-value-for-parameter
