"""
Pipeline de extração de features para treino de modelo próprio.

Lê vídeos de:
  dataset/reais/       → label 0
  dataset/sinteticos/  → label 1
  dataset/cgi/         → label 2

Para cada vídeo extrai 10 frames aleatórios e calcula métricas numéricas.
Salva tudo em dataset_treino.csv.

Uso:
  python ml/extract_features.py
  python ml/extract_features.py --dataset-dir /caminho/para/dataset --output meu_dataset.csv
"""

import argparse
import csv
import os
import random
import sys
from pathlib import Path

import cv2
import numpy as np
import imageio.v3 as iio

# =============================================
# EXTRAÇÃO DE FRAMES
# =============================================

def extract_random_frames(video_path: str, n_frames: int = 10) -> list[np.ndarray]:
    """Extrai N frames aleatórios de um vídeo via imageio (suporta MP4/H.264)."""
    # Conta frames
    total = 0
    for _ in iio.imiter(video_path, plugin="pyav"):
        total += 1

    if total == 0:
        return []

    # Seleciona índices aleatórios ordenados
    n = min(n_frames, total)
    indices = sorted(random.sample(range(total), n))
    target_set = set(indices)

    frames = []
    for i, frame in enumerate(iio.imiter(video_path, plugin="pyav")):
        if i in target_set:
            frames.append(frame)
        if len(frames) >= n:
            break

    return frames


# =============================================
# CÁLCULO DE FEATURES
# =============================================

def compute_features(frame_rgb: np.ndarray) -> dict:
    """
    Calcula todas as features numéricas de um único frame.
    Retorna dict com nomes de coluna → valor.
    """
    gray = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2GRAY)
    features = {}

    # --- 1. PRNU (ruído de sensor) ---
    prnu = _compute_prnu(gray)
    features["prnu_mean"] = prnu["mean"]
    features["prnu_std"] = prnu["std"]
    features["prnu_energy"] = prnu["energy"]

    # --- 2. FFT (picos de alta frequência) ---
    fft = _compute_fft(gray)
    features["fft_high_ratio"] = fft["high_ratio"]
    features["fft_mid_ratio"] = fft["mid_ratio"]
    features["fft_peak_intensity"] = fft["peak_intensity"]
    features["fft_spectral_entropy"] = fft["spectral_entropy"]

    # --- 3. Laplaciano (borrão / nitidez) ---
    lap = _compute_laplacian(gray)
    features["laplacian_var"] = lap["variance"]
    features["laplacian_mean"] = lap["mean"]
    features["laplacian_kurtosis"] = lap["kurtosis"]

    # --- 4. Histograma de Cores (flat colors / CGI) ---
    hist = _compute_color_histogram(frame_rgb)
    features["color_unique_bins"] = hist["unique_bins"]
    features["color_dominant_ratio"] = hist["dominant_ratio"]
    features["color_entropy"] = hist["entropy"]
    features["color_flatness"] = hist["flatness"]

    # --- 5. Textura (variância local) ---
    tex = _compute_texture(gray)
    features["texture_smooth_ratio"] = tex["smooth_ratio"]
    features["texture_var_mean"] = tex["var_mean"]
    features["texture_var_std"] = tex["var_std"]

    # --- 6. Bordas (nitidez de contorno) ---
    edge = _compute_edges(gray)
    features["edge_density"] = edge["density"]
    features["edge_strong_ratio"] = edge["strong_ratio"]

    return features


def _compute_prnu(gray: np.ndarray) -> dict:
    """Estima ruído residual (PRNU) via filtro de Wiener simplificado."""
    # Suaviza para obter a "imagem limpa"
    denoised = cv2.GaussianBlur(gray.astype(np.float64), (5, 5), 0)
    noise = gray.astype(np.float64) - denoised

    return {
        "mean": float(np.mean(np.abs(noise))),
        "std": float(np.std(noise)),
        "energy": float(np.mean(noise ** 2)),
    }


def _compute_fft(gray: np.ndarray) -> dict:
    """Análise de frequência via FFT 2D."""
    f = np.fft.fft2(gray.astype(np.float32))
    f_shift = np.fft.fftshift(f)
    magnitude = np.log1p(np.abs(f_shift))

    h, w = magnitude.shape
    cy, cx = h // 2, w // 2
    max_r = min(cy, cx)

    Y, X = np.ogrid[:h, :w]
    dist = np.sqrt((X - cx) ** 2 + (Y - cy) ** 2)

    low = dist <= (max_r * 0.2)
    mid = (dist > (max_r * 0.2)) & (dist <= (max_r * 0.6))
    high = dist > (max_r * 0.6)

    total_energy = float(np.sum(magnitude)) + 1e-10
    high_energy = float(np.sum(magnitude[high]))
    mid_energy = float(np.sum(magnitude[mid]))

    # Entropia espectral (distribuição de energia)
    mag_flat = magnitude.flatten()
    mag_norm = mag_flat / (mag_flat.sum() + 1e-10)
    mag_norm = mag_norm[mag_norm > 0]
    entropy = float(-np.sum(mag_norm * np.log2(mag_norm + 1e-15)))

    return {
        "high_ratio": high_energy / total_energy,
        "mid_ratio": mid_energy / total_energy,
        "peak_intensity": float(np.max(magnitude[high])) if np.any(high) else 0,
        "spectral_entropy": entropy,
    }


def _compute_laplacian(gray: np.ndarray) -> dict:
    """Variância do Laplaciano — mede nível de borrão."""
    lap = cv2.Laplacian(gray, cv2.CV_64F)
    flat = lap.flatten()

    var = float(np.var(flat))
    mean = float(np.mean(np.abs(flat)))

    # Curtose: distribuição das bordas (IA tende a ser platikúrtica)
    centered = flat - np.mean(flat)
    m4 = float(np.mean(centered ** 4))
    m2 = float(np.var(flat))
    kurtosis = (m4 / (m2 ** 2 + 1e-10)) - 3  # excesso de curtose

    return {
        "variance": var,
        "mean": mean,
        "kurtosis": kurtosis,
    }


def _compute_color_histogram(frame_rgb: np.ndarray) -> dict:
    """Analisa distribuição de cores para detectar flat colors (CGI/anime)."""
    small = cv2.resize(frame_rgb, (160, 90))
    pixels = small.reshape(-1, 3)

    # Quantiza para bins de 16
    quantized = (pixels // 16).astype(np.int32)
    unique_colors = len(set(map(tuple, quantized)))

    # Histograma HSV
    hsv = cv2.cvtColor(small, cv2.COLOR_RGB2HSV)
    hist_h = cv2.calcHist([hsv], [0], None, [32], [0, 180]).flatten()
    hist_s = cv2.calcHist([hsv], [1], None, [32], [0, 256]).flatten()
    hist_v = cv2.calcHist([hsv], [2], None, [32], [0, 256]).flatten()

    combined = np.concatenate([hist_h, hist_s, hist_v])
    combined_norm = combined / (combined.sum() + 1e-10)

    # Entropia de cor
    nonzero = combined_norm[combined_norm > 0]
    entropy = float(-np.sum(nonzero * np.log2(nonzero + 1e-15)))

    # Dominância: fração coberta pelos top 5 bins
    sorted_bins = np.sort(combined_norm)[::-1]
    dominant_ratio = float(np.sum(sorted_bins[:5]))

    # Flatness: proporção de pixels em blocos "flat" (variância local < 10)
    gray = cv2.cvtColor(small, cv2.COLOR_RGB2GRAY).astype(np.float32)
    local_mean = cv2.blur(gray, (8, 8))
    local_sqmean = cv2.blur(gray ** 2, (8, 8))
    local_var = np.maximum(0, local_sqmean - local_mean ** 2)
    flatness = float(np.mean(local_var < 10))

    return {
        "unique_bins": unique_colors,
        "dominant_ratio": dominant_ratio,
        "entropy": entropy,
        "flatness": flatness,
    }


def _compute_texture(gray: np.ndarray) -> dict:
    """Variância local — mede uniformidade de texturas."""
    g = gray.astype(np.float32)
    mean = cv2.blur(g, (5, 5))
    sqmean = cv2.blur(g * g, (5, 5))
    local_var = np.maximum(0, sqmean - mean * mean)

    smooth_ratio = float(np.mean(local_var < 10))

    return {
        "smooth_ratio": smooth_ratio,
        "var_mean": float(np.mean(local_var)),
        "var_std": float(np.std(local_var)),
    }


def _compute_edges(gray: np.ndarray) -> dict:
    """Densidade e qualidade de bordas."""
    all_edges = cv2.Canny(gray, 30, 80)
    strong_edges = cv2.Canny(gray, 100, 200)

    all_density = float(np.mean(all_edges > 0))
    strong_density = float(np.mean(strong_edges > 0))

    return {
        "density": all_density,
        "strong_ratio": strong_density / (all_density + 1e-10),
    }


# =============================================
# PIPELINE PRINCIPAL
# =============================================

FEATURE_COLUMNS = list(compute_features(np.zeros((64, 64, 3), dtype=np.uint8)).keys())

VIDEO_EXTS = {".mp4", ".webm", ".avi", ".mkv", ".mov"}


def process_folder(folder: str, label: int, n_frames: int = 10) -> list[dict]:
    """Processa todos os vídeos de uma pasta, retorna lista de rows."""
    folder_path = Path(folder)
    if not folder_path.exists():
        print(f"  [SKIP] Pasta não encontrada: {folder}")
        return []

    videos = [f for f in folder_path.iterdir() if f.suffix.lower() in VIDEO_EXTS]
    print(f"  Encontrados {len(videos)} vídeos em {folder}")

    rows = []
    for i, video_path in enumerate(videos):
        try:
            frames = extract_random_frames(str(video_path), n_frames)
            if not frames:
                print(f"    [{i+1}/{len(videos)}] {video_path.name}: sem frames, pulando")
                continue

            for j, frame in enumerate(frames):
                features = compute_features(frame)
                features["video_file"] = video_path.name
                features["frame_index"] = j
                features["label"] = label
                rows.append(features)

            print(f"    [{i+1}/{len(videos)}] {video_path.name}: {len(frames)} frames OK")

        except Exception as e:
            print(f"    [{i+1}/{len(videos)}] {video_path.name}: ERRO - {e}")

    return rows


def main():
    parser = argparse.ArgumentParser(description="Extrai features de vídeos para treino ML")
    parser.add_argument("--dataset-dir", default="dataset", help="Pasta raiz do dataset")
    parser.add_argument("--output", default="dataset_treino.csv", help="Arquivo CSV de saída")
    parser.add_argument("--frames", type=int, default=10, help="Frames por vídeo")
    args = parser.parse_args()

    base = Path(args.dataset_dir)

    folders = [
        (base / "reais", 0, "Real"),
        (base / "sinteticos", 1, "IA Generativa"),
        (base / "cgi", 2, "CGI/Animação"),
    ]

    all_rows = []
    for folder, label, name in folders:
        print(f"\n{'='*50}")
        print(f"Processando: {name} (label={label}) — {folder}")
        print(f"{'='*50}")
        rows = process_folder(str(folder), label, args.frames)
        all_rows.extend(rows)
        print(f"  Total: {len(rows)} amostras")

    if not all_rows:
        print("\nNenhuma amostra extraída. Verifique se as pastas existem e contêm vídeos.")
        sys.exit(1)

    # Salva CSV
    fieldnames = ["video_file", "frame_index"] + FEATURE_COLUMNS + ["label"]
    output_path = Path(args.output)

    with open(output_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(all_rows)

    print(f"\n{'='*50}")
    print(f"Dataset salvo: {output_path}")
    print(f"Total de amostras: {len(all_rows)}")

    # Resumo por classe
    from collections import Counter
    counts = Counter(r["label"] for r in all_rows)
    label_names = {0: "Real", 1: "IA", 2: "CGI"}
    for label, count in sorted(counts.items()):
        print(f"  Label {label} ({label_names.get(label, '?')}): {count} amostras")

    print(f"Features por amostra: {len(FEATURE_COLUMNS)}")
    print(f"Colunas: {', '.join(FEATURE_COLUMNS)}")


if __name__ == "__main__":
    main()