#!/usr/bin/env python3
"""
som_auto_adjust_runner.py — V6.7 Standalone executable script for SOM
auto-adjustment with metrics evaluation and failure indicator detection.

User requirement:
    "AVALIAR as métricas de aprendizado e indicadores de falha em Mapas
    Auto-Organizáveis (SOM / Redes de Kohonen) avaliam a fidelidade de
    representação dos dados e a preservação da vizinhança topológica:

    1. Métricas Principais de Aprendizado (gerar scripts capazes de
       autoajustes):
       1.1. Erro de Quantização (QE)
       1.2. Erro Topológico (TE)
       1.3. Erro de Kaski-Lagus
       1.4. Variância Explicada

    2. Indicadores de Não Aprendizado ou Falha (gerar scripts capazes de
       autoajustes):
       2.1. Colapso Topológico
       2.2. Neurônios Mortos (Dead Neurons)
       2.3. Estagnação do QE
       2.4. Cruzamento de Vizinhança

    PORTANTO: observar se os valores demonstram que o modelo esteja
    aprendendo e parar caso não esteja e aprimorar matematicamente a lógica
    de autoajustes analisando Kohonen."

USAGE:
    python3 som_auto_adjust_runner.py --kls-state <path.pt> [--auto-adjust]
    python3 som_auto_adjust_runner.py --live --kls-object <pickle>
    python3 som_auto_adjust_runner.py --demo

OUTPUT:
    /home/z/my-project/download/som_metrics_<timestamp>.json
    /home/z/my-project/download/som_auto_adjust_<timestamp>.json
"""
from __future__ import annotations

import sys
import os
import json
import time
import argparse
import logging
from pathlib import Path
from datetime import datetime
from typing import Any, Dict, List, Optional

# Add project src to path
PROJECT_ROOT = Path("/home/z/my-project/BiGRU_T_version")
sys.path.insert(0, str(PROJECT_ROOT / "src"))

logging.basicConfig(
    level=logging.INFO,
    format="[%(asctime)s] [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("som_auto_adjust_runner")


# ============================================================================
# SECTION 1: SOM METRICS EVALUATION
# ============================================================================

def evaluate_som_metrics(weights: Any, data: Any = None,
                          history: Optional[List[Dict]] = None) -> Dict[str, Any]:
    """Avalia todas as 4 métricas principais + 4 indicadores de falha.

    Métricas Principais (User requirement §1):
        - Quantization Error (QE)
        - Topological Error (TE)
        - Kaski-Lagus Error
        - Explained Variance Share

    Indicadores de Falha (User requirement §2):
        - Topological Collapse
        - Dead Neurons Rate
        - QE Stagnation
        - Neighborhood Crossing

    Args:
        weights: tensor de pesos do SOM (shape: grid_dims + input_dim)
        data: tensor de dados de entrada (opcional)
        history: histórico de métricas anteriores (para detecção de estagnação)

    Returns:
        Dicionário completo com todas as métricas e indicadores.
    """
    import torch
    from bigru_t.model.som_metrics import (
        compute_all_metrics, SOMMetricHistory
    )

    if not isinstance(weights, torch.Tensor):
        weights = torch.tensor(weights, dtype=torch.float32)

    logger.info(f"[METRICS] Avaliando SOM com pesos shape={tuple(weights.shape)}")
    logger.info(f"[METRICS] data={'fornecido' if data is not None else 'None'}")
    logger.info(f"[METRICS] history={len(history) if history else 0} entradas")

    # Se data não for fornecida, usa os próprios pesos como proxy
    if data is None:
        # Flatten weights to use as data samples
        if weights.dim() > 2:
            data = weights.reshape(-1, weights.shape[-1])
        else:
            data = weights
        logger.info(f"[METRICS] data proxy criado: shape={tuple(data.shape)}")

    # Converte history (list of dicts) em SOMMetricHistory (objeto esperado)
    som_history = SOMMetricHistory()
    if history:
        for entry in history:
            som_history.record(entry)
        logger.info(f"[METRICS] SOMMetricHistory populado com {len(history)} entradas")

    try:
        # compute_all_metrics já chama _nan_safe_metrics internamente quando
        # necessário. NÃO devemos chamar _nan_safe_metrics aqui.
        metrics = compute_all_metrics(
            data=data,
            weights=weights,
            history=som_history,
        )

        # Log das métricas principais
        qe = float(metrics.get("quantization_error", 0))
        te = float(metrics.get("topological_error", 0))
        kl = float(metrics.get("kaski_lagus_error", 0))
        ev = float(metrics.get("explained_variance_share", 0))
        logger.info(f"[METRICS] QE (Quantization Error) = {qe:.6f}")
        logger.info(f"[METRICS] TE (Topological Error)  = {te:.6f}")
        logger.info(f"[METRICS] Kaski-Lagus Error       = {kl:.6f}")
        logger.info(f"[METRICS] Explained Variance       = {ev:.6f}")

        # Log dos indicadores de falha
        n_failures = int(metrics.get("n_failure_indicators", 0))
        failures = metrics.get("failure_indicators", [])
        health = metrics.get("overall_health", "unknown")
        logger.info(f"[METRICS] Overall Health: {health}")
        logger.info(f"[METRICS] Failure Indicators: {n_failures} ativos")
        for f in failures:
            logger.warning(f"[METRICS]   ⚠ {f}")

        # Detalhes dos 4 indicadores
        collapse = metrics.get("topological_collapse", {})
        dead = metrics.get("dead_neuron_rate", {})
        stag = metrics.get("qe_stagnation", {})
        cross = metrics.get("neighborhood_crossing", {})
        logger.info(f"[METRICS] Topological Collapse: severity={collapse.get('severity', 'none')}")
        logger.info(f"[METRICS] Dead Neurons: rate={dead.get('dead_neuron_rate', 0):.2%}")
        logger.info(f"[METRICS] QE Stagnation: detected={stag.get('stagnation_detected', False)}")
        logger.info(f"[METRICS] Neighborhood Crossing: severity={cross.get('severity', 'none')}")

        return metrics
    except Exception as e:
        import traceback
        logger.error(f"[METRICS] Erro ao computar métricas: {e}")
        traceback.print_exc()
        return {"error": str(e), "traceback": traceback.format_exc()}


# ============================================================================
# SECTION 2: SOM AUTO-ADJUSTMENT LOGIC
# ============================================================================

def apply_auto_adjustments(metrics: Dict, kls: Optional[Any] = None) -> Dict[str, Any]:
    """Aplica auto-ajustes SOM baseado nas métricas e indicadores de falha.

    User requirement §1.3.1: α₀: 0.5 a 1.0 (decai para ~0.01-0.001)
    User requirement §1.3.2: σ₀: metade da maior dimensão da grade

    Estratégia matemática (Kohonen 2001):
        - Se QE alto (> mediana histórica + 2σ) → aumenta α₀
        - Se TE alto (> 0.2) → aumenta σ₀ (preserva topologia)
        - Se Dead Neurons > 30% → aumenta γ (conscience) e reset σ
        - Se Colapso Topológico → RESET completa dos pesos (severe)
        - Se Estagnação QE → boost de α₀ para escapar de mínimo local
        - Se Cruzamento Vizinhança → reduz α₀ e aumenta σ₀

    Args:
        metrics: dicionário de métricas (saída de evaluate_som_metrics)
        kls: instância de KohonenLearningSystem (opcional — para aplicar mudanças)

    Returns:
        Dicionário com ações tomadas e novos hiperparâmetros.
    """
    from bigru_t.model.som_auto_adjust import create_auto_adjuster

    logger.info("[AUTO-ADJUST] Iniciando análise de auto-ajuste...")
    adjuster = create_auto_adjuster()

    # State before
    state_before = adjuster.state.to_dict() if hasattr(adjuster.state, "to_dict") else {}

    # Aplica análise via auto_adjuster (lógica canônica do projeto)
    actions_taken: List[str] = []
    stop_training = False
    stop_reason: Optional[str] = None

    try:
        # O auto_adjuster consome um dicionário de métricas
        # Signature: adjust(self, kls, som_metrics=None, batch_idx=0)
        # kls pode ser None (apenas analisa, não aplica mudanças)
        result = adjuster.adjust(kls=kls, som_metrics=metrics, batch_idx=0)
        if isinstance(result, dict):
            stop_training = bool(result.get("stop_training", False))
            stop_reason = result.get("stop_reason")
            actions = result.get("actions", result.get("actions_taken", []))
            actions_taken.extend(actions)
            logger.info(f"[AUTO-ADJUST] Ações: {actions}")
            if stop_training:
                logger.warning(f"[AUTO-ADJUST] STOP TRAINING: {stop_reason}")
    except Exception as e:
        logger.error(f"[AUTO-ADJUST] Erro ao aplicar ajuste: {e}")
        actions_taken.append(f"error: {e}")

    # Lógica adicional de auto-ajuste matemático (Kohonen)
    # Baseada nos indicadores de falha específicos
    qe = float(metrics.get("quantization_error", 0))
    te = float(metrics.get("topological_error", 0))
    dead_rate = float(metrics.get("dead_neuron_rate", {}).get("dead_neuron_rate", 0))
    collapse_severity = metrics.get("topological_collapse", {}).get("severity", "none")
    stag_detected = metrics.get("qe_stagnation", {}).get("stagnation_detected", False)
    cross_severity = metrics.get("neighborhood_crossing", {}).get("severity", "none")

    recommendations: Dict[str, Any] = {}

    # Recomendação 1: α₀ baseado em QE
    # Kohonen: α₀ ∈ [0.5, 1.0] para fase de ordenação, decay para ~0.01
    if qe > 0.5:
        recommendations["alpha_0"] = 1.0  # máximo para escapar de mínimos
        recommendations["alpha_decay"] = "exponential"  # decai rápido
        actions_taken.append("alpha_0=1.0 (QE alto, escape de mínimo)")
    elif qe > 0.1:
        recommendations["alpha_0"] = 0.7
        recommendations["alpha_decay"] = "exponential"
        actions_taken.append("alpha_0=0.7 (QE moderado)")
    else:
        recommendations["alpha_0"] = 0.5
        recommendations["alpha_decay"] = "linear"  # decai suave
        actions_taken.append("alpha_0=0.5 (QE baixo, convergência)")

    # Recomendação 2: σ₀ baseado em TE
    # Kohonen: σ₀ = metade da maior dimensão da grade
    # Grid (4,4,4,4) → maior dimensão = 4 → σ₀ = 2
    grid_dims = (4, 4, 4, 4)
    largest_dim = max(grid_dims)
    sigma_0_baseline = largest_dim / 2.0  # = 2.0
    if te > 0.3:
        recommendations["sigma_0"] = sigma_0_baseline * 1.5  # 3.0 — mais vizinhança
        actions_taken.append("sigma_0=3.0 (TE alto, preserva topologia)")
    elif te > 0.1:
        recommendations["sigma_0"] = sigma_0_baseline * 1.25  # 2.5
        actions_taken.append("sigma_0=2.5 (TE moderado)")
    else:
        recommendations["sigma_0"] = sigma_0_baseline  # 2.0
        actions_taken.append("sigma_0=2.0 (TE baixo, canonical)")

    # Recomendação 3: Dead Neurons → conscience mechanism
    if dead_rate > 0.30:
        recommendations["conscience_gamma"] = 0.05  # ativa conscience
        actions_taken.append(f"conscience_gamma=0.05 (dead_rate={dead_rate:.2%})")
    elif dead_rate > 0.10:
        recommendations["conscience_gamma"] = 0.02
        actions_taken.append(f"conscience_gamma=0.02 (dead_rate={dead_rate:.2%})")
    else:
        recommendations["conscience_gamma"] = 0.0  # sem conscience

    # Recomendação 4: Colapso Topológico → RESET
    if collapse_severity == "severe":
        recommendations["reset_weights"] = True
        recommendations["reset_reason"] = "topological_collapse_severe"
        actions_taken.append("RESET pesos (colapso topológico severo)")
        stop_training = True
        stop_reason = "Topological collapse severe — needs full reset"
    elif collapse_severity == "moderate":
        recommendations["boost_alpha"] = 1.5  # multiplica α₀
        actions_taken.append("boost_alpha x1.5 (colapso moderado)")

    # Recomendação 5: Estagnação QE → boost
    if stag_detected:
        recommendations["boost_alpha"] = recommendations.get("boost_alpha", 1.0) * 1.3
        recommendations["sigma_boost"] = 1.2
        actions_taken.append("boost_alpha x1.3 + sigma_boost x1.2 (estagnação QE)")

    # Recomendação 6: Cruzamento de Vizinhança → reduz α, aumenta σ
    if cross_severity in ("severe", "moderate"):
        recommendations["alpha_reduction"] = 0.7
        recommendations["sigma_boost"] = recommendations.get("sigma_boost", 1.0) * 1.4
        actions_taken.append(f"alpha_reduction x0.7 + sigma_boost x1.4 (crossing={cross_severity})")

    # State after
    state_after = adjuster.state.to_dict() if hasattr(adjuster.state, "to_dict") else {}

    result = {
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "metrics_summary": {
            "qe": qe,
            "te": te,
            "dead_rate": dead_rate,
            "collapse_severity": collapse_severity,
            "qe_stagnation": stag_detected,
            "crossing_severity": cross_severity,
        },
        "actions_taken": actions_taken,
        "recommendations": recommendations,
        "stop_training": stop_training,
        "stop_reason": stop_reason,
        "state_before": state_before,
        "state_after": state_after,
    }
    logger.info(f"[AUTO-ADJUST] {len(actions_taken)} ações tomadas")
    if stop_training:
        logger.warning(f"[AUTO-ADJUST] STOP: {stop_reason}")
    return result


# ============================================================================
# SECTION 3: DEMO MODE (test with synthetic SOM)
# ============================================================================

def run_demo() -> Dict[str, Any]:
    """Executa demo com SOM sintético para validar lógica."""
    import torch
    logger.info("[DEMO] Criando SOM sintético (4,4,4,4) = 256 neurons, 4D input")
    # SOM weights: shape (4,4,4,4,4) — last dim is input_dim
    torch.manual_seed(42)
    weights = torch.randn(4, 4, 4, 4, 4) * 0.1  # small init
    # Synthetic data clustered
    data = torch.randn(500, 4) * 0.2
    history: List[Dict] = []
    # Simula 5 épocas de histórico
    for ep in range(5):
        history.append({
            "quantization_error": 0.5 - ep * 0.08,
            "topological_error": 0.2 - ep * 0.03,
            "kaski_lagus_error": 0.4 - ep * 0.06,
            "explained_variance_share": 0.3 + ep * 0.1,
        })
    metrics = evaluate_som_metrics(weights, data, history)
    adjustments = apply_auto_adjustments(metrics)
    return {"metrics": metrics, "adjustments": adjustments}


# ============================================================================
# SECTION 4: MAIN ENTRYPOINT
# ============================================================================

def main() -> int:
    parser = argparse.ArgumentParser(
        description="V6.7 SOM Auto-Adjust Runner — métricas + indicadores de falha"
    )
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--demo", action="store_true",
                      help="Executa demo com SOM sintético")
    mode.add_argument("--kls-state", type=str,
                      help="Caminho para arquivo .pt com estado KLS")
    mode.add_argument("--live", action="store_true",
                      help="Modo live (lê KLS do processo atual — experimental)")
    parser.add_argument("--auto-adjust", action="store_true",
                        help="Aplica auto-ajustes (caso contrário, só avalia)")
    parser.add_argument("--output-dir", type=str,
                        default="/home/z/my-project/download",
                        help="Diretório para salvar resultados JSON")
    args = parser.parse_args()

    print("=" * 70)
    print("V6.7 — SOM AUTO-ADJUST RUNNER (Métricas + Indicadores de Falha)")
    print("=" * 70)
    print("User requirement:")
    print("  Avalia: QE, TE, Kaski-Lagus, Variância Explicada")
    print("  Falhas: Colapso Topológico, Neurônios Mortos,")
    print("          Estagnação QE, Cruzamento Vizinhança")
    print("  Auto-ajuste: α₀ ∈ [0.5, 1.0] (decay ~0.01-0.001),")
    print("               σ₀ = metade da maior dimensão da grade")
    print("=" * 70)

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    ts = datetime.utcnow().strftime("%Y%m%d_%H%M%S")

    if args.demo:
        result = run_demo()
        metrics_file = output_dir / f"som_metrics_demo_{ts}.json"
        adjust_file = output_dir / f"som_auto_adjust_demo_{ts}.json"
        with open(metrics_file, "w") as f:
            json.dump(result["metrics"], f, indent=2, default=str)
        with open(adjust_file, "w") as f:
            json.dump(result["adjustments"], f, indent=2, default=str)
        print(f"\n✓ Métricas salvas: {metrics_file}")
        print(f"✓ Auto-ajustes salvos: {adjust_file}")
        # Print summary
        adj = result["adjustments"]
        print(f"\n--- RESUMO ---")
        print(f"Actions: {len(adj['actions_taken'])}")
        for a in adj["actions_taken"]:
            print(f"  • {a}")
        if adj["stop_training"]:
            print(f"STOP TRAINING: {adj['stop_reason']}")
        return 0

    elif args.kls_state:
        import torch
        logger.info(f"[LOAD] Carregando estado KLS de: {args.kls_state}")
        state = torch.load(args.kls_state, map_location="cpu", weights_only=False)
        # Extrai pesos do SOM do estado
        if isinstance(state, dict):
            # Tenta várias chaves canônicas
            weights = (
                state.get("som_weights") or
                state.get("weights") or
                state.get("som", {}).get("weights") if isinstance(state.get("som"), dict) else
                state.get("som")
            )
        else:
            weights = getattr(state, "weights", None)

        if weights is None:
            logger.error("Não foi possível extrair pesos SOM do estado.")
            return 1

        metrics = evaluate_som_metrics(weights)
        adjustments = apply_auto_adjustments(metrics) if args.auto_adjust else None

        metrics_file = output_dir / f"som_metrics_{ts}.json"
        with open(metrics_file, "w") as f:
            json.dump(metrics, f, indent=2, default=str)
        print(f"\n✓ Métricas salvas: {metrics_file}")

        if adjustments:
            adjust_file = output_dir / f"som_auto_adjust_{ts}.json"
            with open(adjust_file, "w") as f:
                json.dump(adjustments, f, indent=2, default=str)
            print(f"✓ Auto-ajustes salvos: {adjust_file}")
        return 0

    else:
        # --live mode (experimental)
        logger.error("--live mode não implementado nesta versão.")
        return 1


if __name__ == "__main__":
    sys.exit(main())