import numpy as np
import torch
from sklearn.decomposition import PCA, TruncatedSVD

from ..util.waveform_util import (
    channel_subset_by_radius,
    get_channels_in_probe,
    set_channels_in_probe,
)
from ..util.spiketorch import svd_lowrank_helper
from .transform_base import (
    BaseWaveformAutoencoder,
    BaseWaveformDenoiser,
    BaseWaveformFeaturizer,
    BaseWaveformModule,
)


class BaseTemporalPCA(BaseWaveformModule):
    default_name = "basis"

    def __init__(
        self,
        channel_index,
        geom=None,
        rank=8,
        whiten=False,
        centered=False,
        fit_radius=None,
        random_state=0,
        name=None,
        name_prefix="",
        temporal_slice: slice | None = None,
        n_oversamples=10,
        niter=21,
        max_waveforms=20_000,
        batch_size=4096,
        fit_dtype=torch.double,
    ):
        if fit_radius is not None:
            if geom is None or channel_index is None:
                raise ValueError("TemporalPCA with fit_radius!=None requires geom.")

        super().__init__(
            channel_index=channel_index, geom=geom, name=name, name_prefix=name_prefix
        )

        # behavior
        self.rank = rank
        self.centered = centered
        self.whiten = whiten
        self.temporal_slice = temporal_slice
        self.batch_size = batch_size

        # fit control
        self.fit_radius = fit_radius
        self.random_state = random_state
        self.n_oversamples = n_oversamples
        self.niter = niter
        self.fit_dtype = fit_dtype
        self.max_waveforms = max_waveforms

        # gizmo
        self._needs_fit = True
        self.shape = (rank, channel_index.shape[1])
        self.nt = None

    def initialize_spike_length_dependent_params(self):
        nt = self.spike_length_samples
        assert nt is not None
        if self.temporal_slice is None:
            self._temporal_ix = None
        else:
            self.register_buffer("_temporal_ix", torch.arange(nt)[self.temporal_slice])
            nt = self.b._temporal_ix.numel()
        self.nt = nt
        self.register_buffer("mean", torch.zeros(nt))
        self.register_buffer("components", torch.zeros(self.rank, nt))
        self.register_buffer("whitener", torch.zeros(self.rank))
        self.to(self.b.channel_index.device)

    def fit(
        self,
        recording,
        waveforms,
        *,
        channels,
        weights=None,
        time_shifts=None,
        **unused,
    ):
        super().fit(recording, waveforms, channels=channels, weights=weights)
        if weights is not None and waveforms.shape[0] > self.max_waveforms:
            self.random_state = np.random.default_rng(self.random_state)
            weights = weights.numpy(force=True) if torch.is_tensor(weights) else weights
            weights = weights.astype(np.float64)
            weights = weights / weights.sum()
            choices = self.random_state.choice(
                len(weights), p=weights, size=self.max_waveforms
            )
            choices.sort()
            choices = torch.from_numpy(choices)
            waveforms = waveforms[choices]
            channels = channels[choices]
        waveforms = self._temporal_slice(waveforms, time_shifts=time_shifts)
        self.dtype = waveforms.dtype
        if self.fit_radius is not None:
            waveforms, train_channel_index = channel_subset_by_radius(
                waveforms,
                channels,
                self.channel_index.to(device=waveforms.device),
                self.geom.to(device=waveforms.device),
                self.fit_radius,
            )
        else:
            train_channel_index = self.b.channel_index.to(waveforms.device)
        _, waveforms_fit = get_channels_in_probe(
            waveforms, channels, train_channel_index
        )
        waveforms_fit = waveforms_fit.to(self.b.channel_index.device)

        if self.centered:
            mean = waveforms_fit.mean(0)
        else:
            mean = torch.zeros_like(waveforms_fit[0])

        if self.centered:
            # torch does not seem always to want to broadcast M as advertised?
            M = mean[None].broadcast_to(waveforms_fit.shape)
        else:
            M = None

        _, components, _, whitener = svd_lowrank_helper(
            x=waveforms_fit,
            rank=self.rank,
            n_oversamples=self.n_oversamples,
            fit_dtype=self.fit_dtype,
            niter=self.niter,
            M=M,
            with_loadings=False,
        )

        self.b.mean.copy_(mean)
        self.b.components.copy_(components)
        self.b.whitener.copy_(whitener)
        self._needs_fit = False

    def needs_fit(self):
        return self._needs_fit

    def _temporal_slice(
        self, waveforms: torch.Tensor, time_shifts=None
    ) -> torch.Tensor:
        if self.temporal_slice is None:
            return waveforms

        if time_shifts is None:
            return waveforms[:, self.temporal_slice]

        assert self._temporal_ix is not None
        n, t, c = waveforms.shape
        t_ = self._temporal_ix.numel()
        assert t_ <= t
        assert time_shifts.shape == (n,)
        time_ix = self._temporal_ix[None, :, None] + time_shifts[:, None, None]
        waveforms = waveforms.take_along_dim(dim=1, indices=time_ix)
        assert waveforms.shape == (n, t_, c)
        return waveforms

    def _transform_in_probe(self, waveforms_in_probe):
        n = waveforms_in_probe.shape[0]
        if self.whiten:
            W = self.b.components / self.b.whitener
        else:
            W = self.b.components
        out = waveforms_in_probe.new_empty((n, self.rank))
        dev = self.b.components.device
        same_dev = out.device == dev
        for i0 in range(0, n, self.batch_size):
            i1 = min(n, i0 + self.batch_size)
            x = waveforms_in_probe[i0:i1].to(device=dev)
            if self.centered:
                x = x - self.mean
            if same_dev:
                torch.mm(x, W.T, out=out[i0:i1])
            else:
                out[i0:i1] = (x @ W.T).to(out)
        return out

    def _inverse_transform_in_probe(self, features):
        if self.whiten:
            W = self.b.components * self.b.whitener
        else:
            W = self.b.components
        assert self.nt is not None
        n = features.shape[0]
        out = features.new_zeros((n, self.nt))
        dev = self.b.components.device
        for i0 in range(0, n, self.batch_size):
            i1 = min(n, i0 + self.batch_size)
            x = features[i0:i1].to(device=dev)
            if out.device == dev:
                if self.centered:
                    torch.addmm(self.b.mean, x, W, out=out[i0:i1])
                else:
                    torch.mm(x, W, out=out[i0:i1])
            else:
                if self.centered:
                    out[i0:i1] = torch.addmm(self.b.mean, x, W).to(out)
                else:
                    out[i0:i1] = torch.mm(x, W).to(out)
        return out

    def _project_in_probe(self, waveforms_in_probe):
        n = waveforms_in_probe.shape[0]
        out = torch.empty_like(waveforms_in_probe)
        proj = self.b.components.T @ self.b.components
        dev = proj.device
        same_dev = out.device == dev
        for i0 in range(0, n, self.batch_size):
            i1 = min(n, i0 + self.batch_size)
            x = waveforms_in_probe[i0:i1].to(device=dev)
            if same_dev:
                if self.centered:
                    torch.addmm(self.b.mean, x - self.mean, proj, out=out[i0:i1])
                else:
                    torch.mm(waveforms_in_probe, proj, out=out[i0:i1])
            else:
                if self.centered:
                    out[i0:i1] = torch.addmm(self.b.mean, x - self.mean, proj).to(out)
                else:
                    out[i0:i1] = torch.mm(x, proj).to(out)
        return out

    def force_reconstruct(self, features):
        ndim = features.ndim
        if ndim == 2:
            features = features.unsqueeze(0)
        n, r, c = features.shape
        waveforms = features.permute(0, 2, 1).reshape(n * c, r)
        waveforms = self._inverse_transform_in_probe(waveforms)
        waveforms = waveforms.reshape(n, c, -1).permute(0, 2, 1)
        if ndim == 2:
            waveforms = waveforms[0]
        return waveforms

    def force_project(self, features):
        ndim = features.ndim
        if ndim == 2:
            features = features.unsqueeze(0)
        n, t, c = features.shape
        waveforms = features.mT.reshape(n * c, t)
        waveforms = self._project_in_probe(waveforms)
        waveforms = waveforms.reshape(n, c, t).mT
        if ndim == 2:
            waveforms = waveforms[0]
        return waveforms

    def to_sklearn(self) -> PCA:
        pca = PCA(
            n_components=self.rank,
            random_state=self.random_state,  # type: ignore
            whiten=self.whiten,
        )
        pca.mean_ = self.b.mean.numpy(force=True)
        pca.components_ = self.b.components.numpy(force=True)
        pca.explained_variance_ = np.square(self.b.whitener.numpy(force=True))
        pca.temporal_slice = self.temporal_slice  # this is not standard  # type: ignore
        return pca

    @classmethod
    def from_sklearn(cls, channel_index, pca: PCA | TruncatedSVD, temporal_slice=None):
        if isinstance(pca, PCA):
            whiten = pca.whiten  # type: ignore
        elif isinstance(pca, TruncatedSVD):
            whiten = False
        else:
            assert False
        self = cls(
            channel_index,
            rank=pca.n_components,  # type: ignore
            whiten=whiten,
            temporal_slice=temporal_slice,
        )
        self.initialize_from_sklearn(pca)
        return self

    @classmethod
    def convert(cls, other):
        assert not other._needs_fit  # doesn't handle fit_radius right.
        self = cls(
            channel_index=other.channel_index,
            rank=other.rank,
            centered=other.centered,
            whiten=other.whiten,
            temporal_slice=other.temporal_slice,
        )
        self._needs_fit = other._needs_fit
        self.spike_length_samples = other.spike_length_samples
        self.initialize_spike_length_dependent_params()
        self.to(other.channel_index.device)
        self.b.mean.copy_(other.mean)
        self.b.components.copy_(other.components)
        self.b.whitener.copy_(other.whitener)
        return self

    def initialize_from_sklearn(self, pca):
        if self.temporal_slice is None:
            self.spike_length_samples = pca.components_.shape[1]
        else:
            # not really -- this is a hack.
            self.spike_length_samples = (
                self.temporal_slice.stop - self.temporal_slice.start
            )
        self.initialize_spike_length_dependent_params()
        if hasattr(pca, "mean_"):
            self.b.mean.copy_(torch.from_numpy(pca.mean_))
            self.center = not (self.b.mean == 0.0).all()
        else:
            self.b.mean.zero_()
            self.center = False
        self.b.components.copy_(torch.from_numpy(pca.components_))
        self.b.whitener.copy_(torch.from_numpy(pca.explained_variance_)).sqrt_()
        self._needs_fit = False


class TemporalPCADenoiser(BaseWaveformDenoiser, BaseTemporalPCA):
    default_name = "temporal_pca"

    def forward(self, waveforms, *, channels, time_shifts=None, **unused):
        waveforms = self._temporal_slice(waveforms, time_shifts=time_shifts)
        dev = waveforms.device
        channels_in_probe, waveforms_in_probe = get_channels_in_probe(
            waveforms, channels.to(device=dev), self.channel_index.to(device=dev)
        )
        waveforms_in_probe = self._project_in_probe(waveforms_in_probe)
        return set_channels_in_probe(waveforms_in_probe, waveforms, channels_in_probe)


class TemporalPCAFeaturizer(BaseWaveformFeaturizer, BaseTemporalPCA):
    default_name = "tpca_features"

    def transform(
        self,
        waveforms,
        *,
        channels,
        channel_index=None,
        _return_for_combined=False,
        time_shifts=None,
        **unused,
    ):
        waveforms = self._temporal_slice(waveforms, time_shifts=time_shifts)

        if channel_index is None:
            channel_index = self.b.channel_index
        dev = waveforms.device
        channels_in_probe, waveforms_in_probe = get_channels_in_probe(
            waveforms, channels.to(device=dev), channel_index.to(device=dev)
        )
        features_in_probe = self._transform_in_probe(waveforms_in_probe)
        features = features_in_probe.new_full(
            (waveforms.shape[0], self.rank, channel_index.shape[1]), torch.nan
        )
        features = set_channels_in_probe(
            features_in_probe,
            features,
            channels_in_probe,
        )
        if _return_for_combined:
            return (
                waveforms,
                features_in_probe,
                channels_in_probe,
                {self.name: features},
            )

        return {self.name: features}

    def inverse_transform(self, features, channels, channel_index=None):
        if channel_index is None:
            channel_index = self.b.channel_index
        channels_in_probe, features_in_probe = get_channels_in_probe(
            features, channels, channel_index
        )
        reconstructions_in_probe = self._inverse_transform_in_probe(features_in_probe)
        recshp = (features.shape[0], self.b.components.shape[1], channel_index.shape[1])
        reconstructions = torch.full(
            recshp,
            torch.nan,
            dtype=reconstructions_in_probe.dtype,
            device=reconstructions_in_probe.device,
        )
        return set_channels_in_probe(
            reconstructions_in_probe, reconstructions, channels_in_probe
        )


class TemporalPCA(BaseWaveformAutoencoder, TemporalPCAFeaturizer):
    default_name = "tpca_features"

    def forward(self, waveforms, *, channels, time_shifts=None, **unused):
        waveforms, features_in_probe, channels_in_probe, features = self.transform(
            waveforms,
            channels=channels,
            time_shifts=time_shifts,
            _return_for_combined=True,
        )
        reconstructions_in_probe = self._inverse_transform_in_probe(features_in_probe)
        reconstructions = set_channels_in_probe(
            reconstructions_in_probe, waveforms, channels_in_probe
        )
        return reconstructions, features
