import numpy as np
from sklearn.base import BaseEstimator, OutlierMixin
from sklearn.preprocessing import RobustScaler
from sklearn.neighbors import NearestNeighbors, KDTree
from sklearn.decomposition import PCA


class HDS(BaseEstimator, OutlierMixin):
    def __init__(self, scales=(5, 10, 20), contamination=0.12, memory_capacity=500, random_state=42):
        self.scales = scales
        self.contamination = contamination
        self.memory_capacity = memory_capacity
        self.random_state = random_state
        self.memory_keys = None
        self.memory_vals = None
        self.train_scores_dist_ = None

    @staticmethod
    def _rank_score(values):
        values = np.asarray(values, dtype=float)
        order = np.argsort(values, kind="mergesort")
        ranks = np.empty(len(values), dtype=float)
        ranks[order] = np.arange(len(values), dtype=float)
        if len(values) > 1:
            ranks /= len(values) - 1
        return ranks

    def _update_memory_buffer(self, Z, scores):
        if self.memory_keys is None:
            self.memory_keys = Z.copy()
            self.memory_vals = scores.copy()
        else:
            combined_keys = np.vstack([self.memory_keys, Z])
            combined_vals = np.concatenate([self.memory_vals, scores])
            if len(combined_keys) > self.memory_capacity:
                combined_keys = combined_keys[-self.memory_capacity:]
                combined_vals = combined_vals[-self.memory_capacity:]
            self.memory_keys = combined_keys
            self.memory_vals = combined_vals

    def _retrieve_memory_scores(self, Z):
        if self.memory_keys is None:
            return np.zeros(len(Z))
        sim_matrix = np.dot(Z, self.memory_keys.T)
        k_eval = min(5, len(self.memory_keys))
        top_k_indices = np.argsort(sim_matrix, axis=1)[:, -k_eval:]
        retrieved_signals = np.zeros(len(Z))
        for i in range(len(Z)):
            retrieved_signals[i] = np.mean(self.memory_vals[top_k_indices[i]])
        return retrieved_signals

    def _adaptive_token_fusion(self, tokens_matrix, mean_profile=None):
        if mean_profile is None:
            mean_profile = np.mean(tokens_matrix, axis=0, keepdims=True)
        drive = np.exp(-np.abs(tokens_matrix - mean_profile))
        weights = drive / (np.sum(drive, axis=1, keepdims=True) + 1e-8)
        fusion = np.sum(tokens_matrix * weights, axis=1)
        return fusion

    def fit(self, X, y=None):
        X = np.asarray(X, dtype=float)
        if X.ndim != 2:
            raise ValueError("X must be a 2D array.")
        if len(X) < 8:
            raise ValueError("At least 8 samples are required.")
        if not np.isfinite(X).all():
            raise ValueError("X contains NaN or infinite values.")
        if not 0 <= self.contamination < 0.5:
            raise ValueError("contamination must be in [0, 0.5).")

        self._fit_X = X.copy()
        self.scaler_ = RobustScaler()
        Z = self.scaler_.fit_transform(X)
        self.n_samples_in_ = len(Z)
        self.n_features_in_ = Z.shape[1]

        self.scales_ = sorted({max(2, min(int(k), len(Z) - 1)) for k in self.scales})
        sg_list = []
        cov_matrix = np.cov(Z, rowvar=False) + 1e-4 * np.eye(self.n_features_in_)
        inv_cov = np.linalg.pinv(cov_matrix)
        
        for k in self.scales_:
            k_eff = min(k, len(Z) - 1)
            nbrs = NearestNeighbors(n_neighbors=k_eff + 1, metric="euclidean").fit(Z)
            distances, _ = nbrs.kneighbors(Z)
            euc_radius = np.maximum(distances[:, -1], 1e-12)
            mah_radius = np.array([np.sqrt(np.dot(row.T, np.dot(inv_cov, row))) for row in Z])
            hybrid_global = 0.5 * euc_radius + 0.5 * mah_radius
            sg_list.append(hybrid_global)
            
        self.train_sg_raw_ = np.mean(sg_list, axis=0)
        self.train_sg_sorted_ = np.sort(self.train_sg_raw_)

        n_samples = len(Z)
        k_ln = min(max(15, self.n_features_in_ * 3), n_samples - 1)
        k_ln = max(1, min(k_ln, n_samples - 1))
        nbrs_ln = NearestNeighbors(n_neighbors=k_ln + 1, metric="euclidean").fit(Z)
        _, indices_ln = nbrs_ln.kneighbors(Z)

        local_recon_errors = np.zeros(n_samples)
        global_pca = PCA(n_components=max(1, self.n_features_in_ - 1), random_state=self.random_state)
        global_pca.fit(Z)
        global_errors = np.linalg.norm(Z - global_pca.inverse_transform(global_pca.transform(Z)), axis=1)

        for i in range(n_samples):
            neighbors_idx = indices_ln[i]
            local_data = Z[neighbors_idx]
            if len(local_data) <= self.n_features_in_:
                local_recon_errors[i] = global_errors[i]
                continue
            local_center = np.mean(local_data, axis=0)
            centered_local = local_data - local_center
            n_comp_local = max(1, min(self.n_features_in_ - 1, len(local_data) - 2))
            try:
                U, S, Vt = np.linalg.svd(centered_local, full_matrices=False)
                components = Vt[:n_comp_local]
                z_centered = Z[i] - local_center
                z_proj = np.dot(z_centered, components.T)
                z_recon = np.dot(z_proj, components)
                local_recon_errors[i] = np.linalg.norm(z_centered - z_recon)
            except Exception:
                local_recon_errors[i] = global_errors[i]

        hybrid_residuals = 0.5 * local_recon_errors + 0.5 * global_errors
        self.train_residuals_ = hybrid_residuals
        self.train_sln_sorted_ = np.sort(hybrid_residuals)

        k_tree = min(10, n_samples - 1)
        k_tree = max(1, k_tree)
        self.kdtree_ = KDTree(Z, leaf_size=max(10, n_samples // 10))
        node_depths = []
        for z_i in Z:
            dist, _ = self.kdtree_.query(z_i.reshape(1, -1), k=k_tree)
            node_depths.append(dist[0][-1])
        
        self.train_tn_raw_ = np.array(node_depths)
        self.train_stn_sorted_ = np.sort(self.train_tn_raw_)

        train_sg_ranked = self._rank_score(self.train_sg_raw_)
        train_sln_ranked = self._rank_score(self.train_residuals_)
        train_stn_ranked = self._rank_score(self.train_tn_raw_)

        self.train_tokens_ = np.column_stack([train_sg_ranked, train_sln_ranked, train_stn_ranked])
        self.train_mean_profile_ = np.mean(self.train_tokens_, axis=0, keepdims=True)

        fused_scores = self._adaptive_token_fusion(self.train_tokens_, self.train_mean_profile_)
        train_fused_ranked = self._rank_score(fused_scores)

        self._update_memory_buffer(Z, train_fused_ranked)
        memory_context = self._retrieve_memory_scores(Z)
        train_mem_ranked = self._rank_score(memory_context)

        self.train_final_scores_ = self._rank_score(0.85 * train_fused_ranked + 0.15 * train_mem_ranked)
        self.decision_scores_ = self.train_final_scores_

        threshold_index = int(np.floor((1 - self.contamination) * len(self.decision_scores_)))
        threshold_index = min(max(threshold_index, 0), len(self.decision_scores_) - 1)
        self.threshold_ = np.sort(self.decision_scores_)[threshold_index]
        self.labels_ = np.where(self.decision_scores_ >= self.threshold_, -1, 1)

        return self

    def score_samples(self, X):
        X = np.asarray(X, dtype=float)
        Z = self.scaler_.transform(X)
        ref = self.scaler_.transform(self._fit_X)
        n_query = len(Z)
        n_ref = len(ref)

        cov_matrix = np.cov(ref, rowvar=False) + 1e-4 * np.eye(self.n_features_in_)
        inv_cov = np.linalg.pinv(cov_matrix)
        
        sg_list = []
        for k in self.scales_:
            k_eff = min(k, n_ref - 1)
            nbrs = NearestNeighbors(n_neighbors=k_eff + 1, metric="euclidean").fit(ref)
            distances, _ = nbrs.kneighbors(Z)
            euc_radius = np.maximum(distances[:, -1], 1e-12)
            mah_radius = np.array([np.sqrt(np.dot(row.T, np.dot(inv_cov, row))) for row in Z])
            hybrid_global = 0.5 * euc_radius + 0.5 * mah_radius
            
            global_rank = np.searchsorted(self.train_sg_sorted_, hybrid_global) / n_ref
            sg_list.append(global_rank)
        S_G_q = np.mean(sg_list, axis=0)

        k_ln = min(max(15, self.n_features_in_ * 3), n_ref - 1)
        k_ln = max(1, min(k_ln, n_ref - 1))
        nbrs_ln = NearestNeighbors(n_neighbors=k_ln + 1, metric="euclidean").fit(ref)
        _, indices_ln = nbrs_ln.kneighbors(Z)

        global_pca = PCA(n_components=max(1, self.n_features_in_ - 1), random_state=self.random_state).fit(ref)
        global_errors_q = np.linalg.norm(Z - global_pca.inverse_transform(global_pca.transform(Z)), axis=1)

        local_recon_errors_q = np.zeros(n_query)
        for i in range(n_query):
            neighbors_idx = indices_ln[i]
            local_data = ref[neighbors_idx]
            if len(local_data) <= self.n_features_in_:
                local_recon_errors_q[i] = global_errors_q[i]
                continue
            local_center = np.mean(local_data, axis=0)
            centered_local = local_data - local_center
            n_comp_local = max(1, min(self.n_features_in_ - 1, len(local_data) - 2))
            try:
                U, S, Vt = np.linalg.svd(centered_local, full_matrices=False)
                components = Vt[:n_comp_local]
                z_centered = Z[i] - local_center
                z_proj = np.dot(z_centered, components.T)
                z_recon = np.dot(z_proj, components)
                local_recon_errors_q[i] = np.linalg.norm(z_centered - z_recon)
            except Exception:
                local_recon_errors_q[i] = global_errors_q[i]

        hybrid_residuals_q = 0.5 * local_recon_errors_q + 0.5 * global_errors_q
        S_LN_q = np.searchsorted(self.train_sln_sorted_, hybrid_residuals_q) / n_ref

        k_tree = min(10, n_ref - 1)
        k_tree = max(1, k_tree)
        q_dist, _ = self.kdtree_.query(Z, k=k_tree)
        q_depths = q_dist[:, -1]
        S_TN_q = np.searchsorted(self.train_stn_sorted_, q_depths) / n_ref

        tokens_q = np.column_stack([S_G_q, S_LN_q, S_TN_q])
        fused_scores_q = self._adaptive_token_fusion(tokens_q, self.train_mean_profile_)
        fused_ranked_q = self._rank_score(fused_scores_q)

        memory_context_q = self._retrieve_memory_scores(Z)
        memory_ranked_q = self._rank_score(memory_context_q)

        final_query_scores = self._rank_score(0.85 * fused_ranked_q + 0.15 * memory_ranked_q)
        return final_query_scores

    def predict(self, X):
        return np.where(self.score_samples(X) >= self.threshold_, -1, 1)

    def fit_predict(self, X, y=None):
        self.fit(X, y)
        return self.labels_