import pickle
from math import cos, sin
from typing import List, Tuple

import numpy as np
from scipy.sparse import csr_matrix

import anndata as ad
import scanpy as sc
from anndata import AnnData

from moscot._types import ArrayLike
from moscot.problems.space import AlignmentProblem, MappingProblem

ANGLES = [0, 30, 60]


def adata_space_rotate() -> AnnData:
    grid = _make_grid(10)
    adatas = _make_adata(grid, n=3)
    for adata, angle in zip(adatas, ANGLES):
        theta = np.deg2rad(angle)
        rot = np.array([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]])
        adata.obsm["spatial"] = np.dot(adata.obsm["spatial"], rot)

    adata = ad.concat(adatas, label="batch", index_unique="-")
    adata.uns["spatial"] = {}
    return adata


def adata_mapping() -> AnnData:
    grid = _make_grid(10)
    adataref, adata1, adata2 = _make_adata(grid, n=3)
    sc.pp.pca(adataref)

    return ad.concat([adataref, adata1, adata2], label="batch", join="outer", index_unique="-")


def _make_grid(grid_size: int) -> ArrayLike:
    xlimits = ylimits = [0, 10]
    x1s = np.linspace(*xlimits, num=grid_size)  # type: ignore [call-overload]
    x2s = np.linspace(*ylimits, num=grid_size)  # type: ignore [call-overload]
    X1, X2 = np.meshgrid(x1s, x2s)
    return np.vstack([X1.ravel(), X2.ravel()]).T


def _make_adata(grid: ArrayLike, n: int) -> List[AnnData]:
    rng = np.random.default_rng(42)
    X = rng.normal(size=(100, 60))
    return [AnnData(X=csr_matrix(X), obsm={"spatial": grid.copy()}) for _ in range(3)]


def _adata_split(adata: AnnData) -> Tuple[AnnData, AnnData]:
    adataref = adata[adata.obs["batch"] == "0"].copy()
    adataref.obsm.pop("spatial")
    adatasp = adata[adata.obs["batch"] != "0"].copy()
    return adataref, adatasp


def generate_alignment_data() -> None:
    adata = adata_space_rotate()
    ap = AlignmentProblem(adata=adata)
    ap = ap.prepare(batch_key="batch")
    ap = ap.solve(alpha=0.5, epsilon=1)

    with open("alignment_solutions.pkl", "wb") as fname:
        pickle.dump(ap.solutions, fname)


def generate_mapping_data() -> None:
    adata = adata_mapping()
    adataref, adatasp = _adata_split(adata)
    mp = MappingProblem(adataref, adatasp)
    mp = mp.prepare(batch_key="batch", sc_attr={"attr": "X"})
    mp = mp.solve()
    with open("mapping_solutions.pkl", "wb") as fname:
        pickle.dump(mp.solutions, fname)


if __name__ == "__main__":
    generate_alignment_data()
    generate_mapping_data()