import numpy as np
import torch
from torch import nn
from torch.nn import functional as F

class AnalogRouter(nn.Module):

    def __init__(self, dim, fp_dim=2048, hidden=64):
        super().__init__()
        self.dim = int(dim)
        if self.dim < 2 or int(fp_dim) < 1 or int(hidden) < 1:
            raise ValueError('Invalid spectrum/fingerprint/hidden dimension')
        self.encoder = nn.Sequential(nn.Linear(int(fp_dim), int(hidden)), nn.SiLU())
        self.context = nn.Sequential(nn.Linear(3 * int(hidden) + 4, int(hidden)), nn.SiLU())
        self.router = nn.Sequential(nn.Linear(int(hidden) + 3, 32), nn.SiLU(), nn.Linear(32, 3))
        self.register_buffer('mz', torch.arange(self.dim, dtype=torch.float32), persistent=False)

    def forward(self, fpq, fpr, ref_l1, massq, massr, sim, mode='transport', return_valid=False):
        if mode not in ('transport', 'reweight'):
            raise ValueError('Unknown router mode: ' + str(mode))
        if ref_l1.ndim != 2 or ref_l1.shape[1] != self.dim:
            raise ValueError('ref_l1 must have shape [batch, spectrum dimension]')
        batch = ref_l1.shape[0]
        massq, massr, sim = (massq.reshape(batch), massr.reshape(batch), sim.reshape(batch))
        eq, er = (self.encoder(fpq), self.encoder(fpr))
        globals_ = torch.stack((sim, massq / self.dim, massr / self.dim, (massq - massr) / self.dim), dim=1)
        context = self.context(torch.cat((eq, er, eq - er, globals_), dim=1))
        mz = self.mz.to(dtype=ref_l1.dtype).expand(batch, -1)
        ref = ref_l1.clamp_min(0)
        ref = ref / ref.sum(1, keepdim=True).clamp_min(1e-12)
        peak_features = torch.stack((mz / self.dim, (massr[:, None] - mz) / self.dim, torch.sqrt(ref.clamp_min(0))), dim=2)
        logits = self.router(torch.cat((context[:, None, :].expand(-1, self.dim, -1), peak_features), dim=2))
        log_weights = F.log_softmax(logits, dim=2)
        original = self.mz.to(torch.long).expand(batch, -1)
        delta = torch.round(massq).to(torch.long) - torch.round(massr).to(torch.long)
        shifted = original + delta[:, None] if mode == 'transport' else original
        upper = torch.minimum(torch.round(massq).to(torch.long) + 3, torch.full_like(delta, self.dim - 1))
        source_valid = original >= 1
        copy_valid = source_valid & (original <= upper[:, None])
        shift_valid = source_valid & (shifted >= 1) & (shifted <= upper[:, None]) & (shifted < self.dim)
        positive = ref > 0
        safe_ref = torch.where(positive, ref, torch.ones_like(ref))
        log_ref = torch.log(safe_ref)
        retained = torch.cat((copy_valid & positive, shift_valid & positive), dim=1)
        log_mass = torch.cat((log_ref + log_weights[:, :, 0], log_ref + log_weights[:, :, 1]), dim=1)
        log_mass = log_mass.masked_fill(~retained, -float('inf'))
        valid = retained.any(dim=1)
        safe_log_mass = torch.where(valid[:, None], log_mass, torch.zeros_like(log_mass))
        centered_log_mass = safe_log_mass - safe_log_mass.amax(dim=1, keepdim=True)
        log_total = torch.logsumexp(centered_log_mass, dim=1, keepdim=True)
        normalized_mass = torch.where(retained, torch.exp(centered_log_mass - log_total), torch.zeros_like(safe_log_mass))
        copy_mass, shift_mass = normalized_mass.split(self.dim, dim=1)
        output = torch.zeros_like(ref)
        output.scatter_add_(1, original, copy_mass)
        output.scatter_add_(1, shifted.clamp(0, self.dim - 1), shift_mass)
        return (output, valid) if return_valid else output

def _router_data(data):
    required = ('features', 'masses', 'spectra', 'neighbors', 'similarities')
    missing = [key for key in required if key not in data]
    if missing:
        raise ValueError('Missing router data fields: ' + ', '.join(missing))
    features = np.asarray(data['features'], dtype=np.float32)
    masses = np.asarray(data['masses'], dtype=np.float32)
    spectra = np.asarray(data['spectra'], dtype=np.float32)
    neighbors = np.asarray(data['neighbors'], dtype=np.int64)
    similarities = np.asarray(data['similarities'], dtype=np.float32)
    n = len(features)
    if features.ndim != 2 or features.shape[1] < 1 or masses.shape != (n,):
        raise ValueError('Invalid fingerprint/mass shapes')
    if spectra.ndim != 2 or spectra.shape[0] != n or spectra.shape[1] < 2:
        raise ValueError('Invalid spectrum shape')
    if neighbors.ndim != 2 or neighbors.shape[0] != n or similarities.shape != neighbors.shape:
        raise ValueError('Invalid neighbour/similarity shapes')
    if not all((np.isfinite(x).all() for x in (features, masses, spectra, similarities))):
        raise ValueError('Router data contain non-finite values')
    if (spectra < 0).any() or (masses <= 0).any():
        raise ValueError('Spectra must be nonnegative and masses positive')
    if (neighbors < -1).any() or (neighbors >= n).any():
        raise ValueError('Neighbour indices must be -1 or valid molecular indices')
    if (similarities < 0).any() or (similarities > 1 + 1e-06).any():
        raise ValueError('Similarities must lie in [0, 1]')
    return dict(features=features, masses=masses, spectra=spectra, neighbors=neighbors, similarities=similarities)

def _router_indices(indices, n, label):
    x = np.asarray(indices, dtype=np.int64).reshape(-1)
    if (x < 0).any() or (x >= n).any() or len(np.unique(x)) != len(x):
        raise ValueError(label + ' must contain distinct valid molecular indices')
    return x

def _router_eligible(data, indices, min_similarity):
    refs = data['neighbors'][indices]
    sims = data['similarities'][indices]
    eligible = (refs >= 0) & (sims >= float(min_similarity))
    eligible &= refs != indices[:, None]
    safe = np.maximum(refs, 0)
    eligible &= data['spectra'][safe].sum(2) > 1e-12
    return eligible

def _router_forward_pairs(model, data, queries, refs, sims, device, mode):
    tensor = lambda x: torch.as_tensor(x, dtype=torch.float32, device=device)
    return model(tensor(data['features'][queries]), tensor(data['features'][refs]), tensor(data['spectra'][refs]), tensor(data['masses'][queries]), tensor(data['masses'][refs]), tensor(sims), mode=mode, return_valid=True)

def predict_router(model, data, query_indices, device='cpu', mode='transport', batch_size=64, min_similarity=0.35):
    data = _router_data(data)
    queries = _router_indices(query_indices, len(data['features']), 'query_indices')
    if int(batch_size) < 1:
        raise ValueError('batch_size must be positive')
    eligible = _router_eligible(data, queries, min_similarity)
    where_q, where_k = np.nonzero(eligible)
    result = np.zeros((len(queries), data['spectra'].shape[1]), dtype=np.float32)
    weight_sum = np.zeros(len(queries), dtype=np.float64)
    previous_training = model.training
    model.to(device)
    model.eval()
    try:
        with torch.no_grad():
            for start in range(0, len(where_q), int(batch_size)):
                positions = where_q[start:start + int(batch_size)]
                columns = where_k[start:start + int(batch_size)]
                q = queries[positions]
                r = data['neighbors'][q, columns]
                sims = data['similarities'][q, columns]
                p, valid = _router_forward_pairs(model, data, q, r, sims, device, mode)
                p = p.cpu().numpy()
                weights = np.power(sims.astype(np.float64), 4) * valid.cpu().numpy()
                np.add.at(result, positions, p * weights[:, None].astype(np.float32))
                np.add.at(weight_sum, positions, weights)
    finally:
        model.train(previous_training)
    covered = weight_sum > 1e-12
    result[covered] /= weight_sum[covered, None].astype(np.float32)
    totals = result.sum(1, dtype=np.float64)
    covered &= totals > 1e-12
    result[covered] /= totals[covered, None].astype(np.float32)
    result[~covered] = 0
    return (result, covered)