from itertools import cycle

import numpy as np
from _petsc_miner import PatternMiner
from aeon.transformations.base import BaseTransformer

from .preprocessing import SAX


class PetsTransformer(BaseTransformer):
    """PETS: Create Pattern-based Embeddings from Time Series."""

    _tags = {
        "X_inner_type": ["numpy3D", "np-list"],
        "y_inner_type": "numpy1D",
        "capability:unequal_length": True,
        "capability:multivariate": True,
        "capability:predict_proba": True,
        "output_data_type": "Tabular",
        "fit_is_empty": False,
        "requires_y": False,
    }

    def __init__(
        self,
        window,
        stride,
        w,
        alpha,
        min_size,
        max_size,
        duration,
        k,
        sort_alpha,
        multiresolution,
        soft,
        tau,
        verbosity,
        cache,
    ):
        self.min_size = min_size
        self.w = w
        max_size = max_size if max_size is not None else w
        self.max_size = max_size
        self.multiresolution = multiresolution
        self.soft = soft
        self.tau = tau if tau is not None else 1 / (2 * alpha)
        self.verbosity = verbosity
        self.cache = cache

        self.sax = SAX(window, stride, w, alpha)
        self.miner = PatternMiner(alpha, min_size, max_size, duration, k, sort_alpha)

        self.patterns_ = []
        self.windows_ = []
        self._data = []

    def fit(self, X, y=None):
        # For MR-PETSC, maximum window size is length of shortest time series in X
        limit = min(row.shape[1] for row in X)
        # Mine each multivariate signal separately
        for dimension, ts in enumerate(self._get_signals(X)):
            # Start with window at min ts length and reduce window each iteration
            if self.multiresolution:
                # Check if w complies with limit
                if self.w > limit:
                    raise ValueError(
                        f"w can be at most the size of the shortest time series in the collection. {self.w=}, {limit=}."
                    )
                self.sax.window = self.w

            while True:
                if self.verbosity >= 1:
                    print(f"window={self.sax.window}")
                    print(f"Dimension {dimension}: applying SAX transform...")
                discrete, labels = self.sax.transform(ts)
                if self.verbosity >= 1:
                    print(f"Dimension {dimension} mining patterns...")
                patterns = self.miner.mine(discrete)
                if self.verbosity >= 1:
                    print(f"Dimension {dimension}: mined {len(patterns)} patterns!")

                self.patterns_.append(patterns)
                self.windows_.append(self.sax.window)

                # Maybe store for use in _transform
                if self.cache:
                    self._data.append((discrete, labels))

                # Add more patterns with halved window while window >= SAX word size
                if self.multiresolution:
                    self.sax.window *= 2
                    if self.sax.window > limit:
                        break
                else:
                    break
        return self

    def fit_transform(self, X, y=None):
        self.fit(X)
        if self.cache:
            return self._transform_cached()
        return self._transform(X, train=True)

    def transform(self, X):
        return self._transform(X, train=False)

    def _transform_cached(self):
        embedding = []
        for (discrete, labels), patterns in zip(self._data, self.patterns_):
            col = self._get_col(discrete, labels, patterns, train=True)
            embedding.append(col)
        embedding = np.concatenate(embedding, axis=1)
        return embedding

    def _transform(self, X, train):
        embedding = []
        it = zip(cycle(self._get_signals(X)), self.windows_, self.patterns_)
        for ts, window, patterns in it:
            self.sax.window = window
            discrete, labels = self.sax.transform(ts)
            col = self._get_col(discrete, labels, patterns, train)
            embedding.append(col)
        embedding = np.concatenate(embedding, axis=1)
        return embedding

    def _get_col(self, discrete, labels, patterns, train):
        col = np.zeros((len(set(labels)), len(patterns)), dtype=int)
        for pat_idx, pattern in enumerate(patterns):
            if train:
                projection = pattern.projection
            elif self.soft:
                projection = self.miner.project_soft(discrete, pattern, self.tau)
            else:
                projection = self.miner.project(discrete, pattern)
            if len(projection) > 0:
                # Sum the number of windows in each time series where pattern occurs
                ts_indices = np.array([labels[idx] for idx in projection])
                np.add.at(col, (ts_indices, pat_idx), 1)
        return col

    def _get_signals(self, X):
        """Split X into list of separate multivariate signals."""
        return [[x[signal] for x in X] for signal in range(X[0].shape[0])]
