import random
from functools import partial
from pathlib import Path

import torch

from .base import BatchMatrixData, GeneratorConfig, MatrixDomain, MatrixGenerator
from .domains import (
    DirectedPowerLawGenerator,
    DiffusionAdvectionGenerator,
    DiffusionGenerator,
    ElasticityGenerator,
    EnhancedAdvectionGenerator,
    EnhancedDiffusionGenerator,
    FastGraphLaplacianGenerator,
    RandomSparseGenerator,
    SBMGraphGenerator,
    SpectralStressGenerator,
    StokesSaddlePointGenerator,
    VariableAdvectionGenerator,
    VariableGridDiffusionGenerator,
)
from .pool import MultiSizePooledGenerator, TopologyPool


class MatrixGeneratorRegistry:

    def __init__(self, generators: dict[MatrixDomain, MatrixGenerator]):
        self.generators = generators

    def generate(self, domain: MatrixDomain, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        generator = self.generators[domain]
        return generator.generate_batch(batch_size, num_context_pairs)

    def sample_domain(self, weights: dict[MatrixDomain, float] | None = None) -> MatrixDomain:
        domains = list(self.generators.keys())
        if weights is None:
            return random.choice(domains)

        domain_list = [d for d in domains if d in weights]
        weight_values = [weights[d] for d in domain_list]
        return random.choices(domain_list, weights=weight_values, k=1)[0]

    @property
    def domains(self) -> list[MatrixDomain]:
        return list(self.generators.keys())


def _build_multi_size_pool(
    generator,
    ns: tuple[int, ...],
    pool_size: int,
    num_workers: int,
    **topology_kwargs,
) -> MultiSizePooledGenerator:
    per_size_pool = max(1, pool_size // len(ns))
    pools = {}
    for n in ns:
        gen_fn = partial(generator.generate_topology_cpu, n, **topology_kwargs)
        pools[n] = TopologyPool.build(gen_fn, per_size_pool, num_workers)
    return MultiSizePooledGenerator(pools, generator.fill_values_gpu)


def build_training_registry(
    config: GeneratorConfig,
    device: torch.device,
    graph_pool_size: int = 1000,
    graph_pool_workers: int = 1,
) -> MatrixGeneratorRegistry:
    grid_sizes = config.grid_sizes
    ns = tuple(gs * gs for gs in grid_sizes)

    diff_gen = DiffusionGenerator(grid_sizes, device)
    graph_laplacian_gen = FastGraphLaplacianGenerator(ns, device)
    sbm_gen = SBMGraphGenerator(ns, device)
    directed_pl_gen = DirectedPowerLawGenerator(ns, device)

    generators: dict[MatrixDomain, MatrixGenerator] = {
        MatrixDomain.DIFFUSION: diff_gen,
        MatrixDomain.DIFFUSION_ADVECTION: DiffusionAdvectionGenerator(diff_gen),
        MatrixDomain.GRAPH_LAPLACIAN: _build_multi_size_pool(
            graph_laplacian_gen, ns, graph_pool_size, graph_pool_workers, m=3,
        ),
        MatrixDomain.ELASTICITY: ElasticityGenerator(
            tuple(gs // 2 for gs in grid_sizes), device,
        ),
        MatrixDomain.STOKES: StokesSaddlePointGenerator(
            tuple(gs // 2 for gs in grid_sizes), device,
        ),
        MatrixDomain.SBM: _build_multi_size_pool(
            sbm_gen, ns, graph_pool_size, graph_pool_workers,
        ),
        MatrixDomain.SPECTRAL_STRESS: SpectralStressGenerator(diff_gen, device),
        MatrixDomain.VARIABLE_DIFFUSION: VariableGridDiffusionGenerator(config, device),
        MatrixDomain.VARIABLE_ADVECTION: VariableAdvectionGenerator(config, device),
        MatrixDomain.ENHANCED_DIFFUSION: EnhancedDiffusionGenerator(config, device),
        MatrixDomain.ENHANCED_ADVECTION: EnhancedAdvectionGenerator(config, device),
        MatrixDomain.RANDOM_SPARSE: RandomSparseGenerator(ns, device),
        MatrixDomain.DIRECTED_POWER_LAW: _build_multi_size_pool(
            directed_pl_gen, ns, graph_pool_size, graph_pool_workers, m=3,
        ),
    }

    return MatrixGeneratorRegistry(generators)


def build_eval_registry(
    config: GeneratorConfig,
    device: torch.device,
    suitesparse_dir: Path | None = None,
) -> MatrixGeneratorRegistry:
    registry = build_training_registry(config, device)

    if suitesparse_dir is not None:
        from .domains.suitesparse_loader import SuiteSparseLoader
        registry.generators[MatrixDomain.SUITESPARSE] = SuiteSparseLoader(suitesparse_dir, device)

    return registry
