from typing import cast, Self

import h5py
import numpy as np
import torch

from ..util.internal_config import ClusteringFeaturesConfig
from ..util import drift_util, interpolation_util
from ..util.py_util import databag
from ..util.waveform_util import single_channel_index
from . import cluster_util

default_clustering_features_cfg = ClusteringFeaturesConfig()
minimal_features_cfg = ClusteringFeaturesConfig(
    n_main_channel_pcs=0,
    use_amplitude=False,
    use_signed_amplitude=False,
    use_x=False,
    use_z=False,
)


def get_clustering_features(
    recording,
    sorting,
    motion_est=None,
    clustering_features_cfg: (
        ClusteringFeaturesConfig | None
    ) = default_clustering_features_cfg,
) -> "SimpleMatrixFeatures":
    if clustering_features_cfg is None:
        clustering_features_cfg = minimal_features_cfg
    if clustering_features_cfg.features_type == "simple_matrix":
        return SimpleMatrixFeatures.from_config(
            recording, sorting, motion_est, clustering_features_cfg
        )
    assert False


@databag
class SimpleMatrixFeatures:
    n: int
    features: np.ndarray
    x: np.ndarray
    z: np.ndarray
    z_reg: np.ndarray
    xyza: np.ndarray
    signed_amplitudes: np.ndarray
    amplitudes: np.ndarray

    def mask(self, mask) -> Self:
        x = self.x[mask]
        return self.__class__(
            n=x.shape[0],
            features=self.features[mask],
            x=x,
            z=self.z[mask],
            z_reg=self.z_reg[mask],
            xyza=self.xyza[mask],
            signed_amplitudes=self.signed_amplitudes[mask],
            amplitudes=self.amplitudes[mask],
        )

    @classmethod
    def from_config(cls, recording, sorting, motion_est, clustering_features_cfg):
        assert clustering_features_cfg.features_type == "simple_matrix"

        xyza = getattr(sorting, clustering_features_cfg.localizations_dataset_name)
        x = xyza[:, 0]
        z_reg = z = xyza[:, 2]
        if motion_est is not None:
            z_reg = motion_est.correct_s(sorting.times_seconds, z)

        features = []

        if clustering_features_cfg.use_z:
            if clustering_features_cfg.motion_aware:
                features.append(z_reg[:, None])
            else:
                features.append(z[:, None])

        if clustering_features_cfg.use_x:
            features.append(x[:, None] * clustering_features_cfg.x_scale)

        amp = getattr(sorting, clustering_features_cfg.amplitudes_dataset_name)
        if clustering_features_cfg.use_amplitude:
            assert amp is not None
            ampft = amp.copy()
            if clustering_features_cfg.log_transform_amplitude:
                ampft = np.log(clustering_features_cfg.amp_log_c + ampft)
                ampft *= clustering_features_cfg.amp_scale
            features.append(ampft[:, None])

        v = getattr(sorting, clustering_features_cfg.voltages_dataset_name, None)
        if v is None:
            samp = amp.copy()
        else:
            samp = amp * np.sign(v)

        if clustering_features_cfg.use_signed_amplitude:
            samp *= clustering_features_cfg.amp_scale
            features.append(clustering_features_cfg.amp_scale * samp[:, None])

        do_pcs = bool(clustering_features_cfg.n_main_channel_pcs)
        pcs = None
        if do_pcs and not clustering_features_cfg.motion_aware:
            pcs = cluster_util.get_main_channel_pcs(
                sorting,
                rank=clustering_features_cfg.n_main_channel_pcs,
                dataset_name=clustering_features_cfg.pca_dataset_name,
            )
        elif do_pcs and clustering_features_cfg.motion_aware:
            geom = sorting.geom
            if motion_est is None:
                registered_geom = geom
            else:
                registered_geom = drift_util.registered_geometry(
                    geom, motion_est=motion_est
                )
            res = drift_util.get_shift_info(sorting, motion_est, geom)
            channels, shifts, n_pitches_shift = res
            mainchan_ci = single_channel_index(len(geom))
            schan, *_ = drift_util.get_stable_channels(
                geom,
                channels,
                mainchan_ci,
                registered_geom,
                n_pitches_shift,
                workers=clustering_features_cfg.workers,
            )
            mask = np.ones((1,), dtype=bool)
            mask = np.broadcast_to(mask, len(schan))
            with h5py.File(sorting.parent_h5_path, "r", locking=False) as h5:
                pcs = interpolation_util.interpolate_by_chunk(
                    mask=mask,
                    dataset=h5[clustering_features_cfg.pca_dataset_name],
                    geom=geom,
                    channel_index=cast(h5py.Dataset, h5["channel_index"])[:],
                    channels=sorting.channels,
                    shifts=shifts,
                    registered_geom=registered_geom,
                    target_channels=schan,
                    params=clustering_features_cfg.interp_params,
                )
                pcs = pcs[:, : clustering_features_cfg.n_main_channel_pcs, 0]

        if do_pcs:
            assert pcs is not None
            if clustering_features_cfg.pc_transform == "log":
                pcs = signed_log1p(
                    pcs, pre_scale=clustering_features_cfg.pc_pre_transform_scale
                )
            elif clustering_features_cfg.pc_transform == "sqrt":
                pcs = signed_sqrt_transform(
                    pcs, pre_scale=clustering_features_cfg.pc_pre_transform_scale
                )
            else:
                assert clustering_features_cfg.pc_transform in ("none", None)
            pcs *= clustering_features_cfg.pc_scale
            if torch.is_tensor(pcs):
                pcs = pcs.numpy(force=True)
            features.append(pcs)

        n = x.shape[0]
        if len(features):
            features = np.concatenate(features, axis=1)
        else:
            features = np.empty((n, 0))
        return cls(
            n=n,
            features=features,
            x=x,
            z=z,
            z_reg=z_reg,
            xyza=xyza,
            amplitudes=amp,
            signed_amplitudes=samp,
        )


def signed_log1p(x, pre_scale=1.0):
    """sgn(x) * log(1+|x|*pre_scale)"""
    x = torch.asarray(x)
    xx = x.abs()
    if pre_scale != 1.0:
        xx.mul_(pre_scale)
    torch.log1p(xx, out=xx)
    xx.mul_(torch.sign(x))
    return xx


def signed_sqrt_transform(x, pre_scale=1.0):
    """sgn(x) * (sqrt(1 + |x|*pre_scale) - 1)"""
    x = torch.asarray(x)
    xx = x.abs()
    if pre_scale != 1.0:
        xx.mul_(pre_scale)
    xx.add_(1.0)
    xx.sqrt_()
    xx.sub_(1.0)
    xx.mul_(torch.sign(x))
    return xx
