import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import tensorflow as tf
import gc

from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix

from tensorflow.keras.datasets import cifar10
from tensorflow.keras import layers, models, callbacks
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.applications import ResNet50V2

# ================================================
# CLEAR MEMORY BEFORE STARTING
# ================================================
tf.keras.backend.clear_session()
gc.collect()
print("Memory cleared.")

# ================================================
# CONFIGURATION
# ================================================
RANDOM_SEED          = 42
np.random.seed(RANDOM_SEED)
tf.random.set_seed(RANDOM_SEED)

VAL_SIZE             = 5000
VARIANCE_THRESHOLD   = 0.95
SVD_ENERGY_THRESHOLD = 0.95
NUM_PER_CLASS        = 50        # 50 per class → 500 distilled total
BLOCK_SIZE           = 3
THETA                = np.pi * (5/3)   # 300°
IMG_SIZE             = 160
TTA_STEPS            = 10
EPOCHS               = 400
BATCH_SIZE           = 8

CIFAR_CLASSES = [
    "airplane", "automobile", "bird", "cat", "deer",
    "dog", "frog", "horse", "ship", "truck"
]

print(f"Rotation angle   : θ = {np.degrees(THETA):.1f}°")

# ================================================
# 1. LOAD CIFAR-10 — ORIGINAL SPLIT
# ================================================
print("=" * 60)
print("1. Loading CIFAR-10 with original train/test split...")
print("=" * 60)

(X_train_raw, y_train_raw), (X_test_raw, y_test_raw) = cifar10.load_data()

y_train = y_train_raw.flatten()
y_test  = y_test_raw.flatten()

print(f"Train pool : {X_train_raw.shape[0]} images")
print(f"Test  set  : {X_test_raw.shape[0]}  images ({X_test_raw.shape[0]//10}/class)")

# ================================================
# 2. PREPROCESSING
# ================================================
print("\n" + "=" * 60)
print("2. Preprocessing...")
print("=" * 60)

X_train_flat = X_train_raw.astype(np.float32).reshape(X_train_raw.shape[0], -1) / 255.0
X_test_flat  = X_test_raw.astype(np.float32).reshape(X_test_raw.shape[0],  -1) / 255.0

X_mean = np.mean(X_train_flat, axis=0)
X_std  = np.std(X_train_flat,  axis=0)
X_std[X_std == 0] = 1e-8

X_train_std = (X_train_flat - X_mean) / X_std
X_test_std  = (X_test_flat  - X_mean) / X_std

X_train_img = X_train_raw.astype(np.float32) / 255.0
X_test_img  = X_test_raw.astype(np.float32)  / 255.0

print(f"Standardized flat : {X_train_std.shape}")
print(f"Raw float images  : {X_train_img.shape}")

# ================================================
# 3. CARVE VALIDATION SET
# ================================================
print("\n" + "=" * 60)
print("3. Carving validation set...")
print("=" * 60)

train_idx = np.arange(X_train_raw.shape[0])
dist_idx, val_idx = train_test_split(
    train_idx,
    test_size=VAL_SIZE,
    stratify=y_train,
    random_state=RANDOM_SEED
)

X_val_img  = X_train_img[val_idx]
y_val      = y_train[val_idx]
X_dist_std = X_train_std[dist_idx]
y_dist_raw = y_train[dist_idx]

print(f"Distillation pool : {X_dist_std.shape[0]} images")
print(f"Validation set    : {X_val_img.shape[0]} images ({X_val_img.shape[0]//10}/class)")
print(f"Test set          : {X_test_img.shape[0]} images ({X_test_img.shape[0]//10}/class)")

# ================================================
# 4. PCA DISTILLATION
# ================================================
print("\n" + "=" * 60)
print("4. Per-class PCA Distillation (95% variance, 4500/class)...")
print("=" * 60)

def pca_distillation(X, y,
                     n_per_class=4500,
                     var_threshold=VARIANCE_THRESHOLD,
                     random_seed=RANDOM_SEED):
    classes     = np.unique(y)
    distilled_X = []
    distilled_y = []
    plot_done   = False

    print(f"\n{'Class':<6} {'Name':<14} {'#PCs':<8} {'Var%':<10} {'Pool'}")
    print("-" * 50)

    for cls in classes:
        mask  = (y == cls)
        X_cls = X[mask]

        pca = PCA(n_components=var_threshold, random_state=random_seed)
        pca.fit(X_cls)
        n_components  = pca.n_components_   # FIX: added underscore
        var_explained = pca.explained_variance_ratio_.sum() * 100
        scores        = pca.transform(X_cls)
        n_select      = min(n_per_class, X_cls.shape[0])

        scores_selected = scores[:n_select, :]
        vectors_recon   = pca.inverse_transform(scores_selected)

        distilled_X.append(vectors_recon)
        distilled_y.append(np.full(n_select, cls))

        print(f"{cls:<6} {CIFAR_CLASSES[cls]:<14} {n_components:<8} "
              f"{var_explained:<10.2f} {n_select}")

        if not plot_done:
            cumvar = np.cumsum(pca.explained_variance_ratio_)
            plt.figure(figsize=(8, 4))
            plt.plot(range(1, len(cumvar) + 1), cumvar, linewidth=1.5)
            plt.axhline(y=var_threshold, color='r', linestyle='--',
                        label=f'{var_threshold*100:.1f}% threshold')
            plt.xlabel('Number of Principal Components')
            plt.ylabel('Cumulative Explained Variance')
            plt.title(f'PCA Variance — Class 0 ({CIFAR_CLASSES[0]})')
            plt.legend(); plt.tight_layout()
            plt.savefig('pca_variance_plot.png', dpi=150)
            plt.show()
            plot_done = True

    distilled_X = np.vstack(distilled_X)
    distilled_y = np.concatenate(distilled_y)
    print(f"\nPCA distillation complete -> {distilled_X.shape[0]} vectors (SVD input)")
    return distilled_X, distilled_y


X_pca_distilled, y_pca_distilled = pca_distillation(X_dist_std, y_dist_raw)

# ================================================
# 5. SVD BLOCK ROTATION HELPERS
# ================================================

def build_rotation_matrix(theta, size=3):
    R = np.eye(size)
    R[0, 0] =  np.cos(theta)
    R[0, 1] = -np.sin(theta)
    R[1, 0] =  np.sin(theta)
    R[1, 1] =  np.cos(theta)
    return R


def block_rotate(Lm, theta, block_size=BLOCK_SIZE):
    m, r  = Lm.shape
    R_rot = build_rotation_matrix(theta, size=block_size)
    remainder = m % block_size
    if remainder != 0:
        pad_rows  = block_size - remainder
        Lm_padded = np.vstack([Lm, np.zeros((pad_rows, r), dtype=Lm.dtype)])
    else:
        Lm_padded = Lm.copy()
    m_padded          = Lm_padded.shape[0]
    Lm_rotated_padded = np.zeros_like(Lm_padded)
    for start in range(0, m_padded, block_size):
        block = Lm_padded[start:start + block_size, :]
        Lm_rotated_padded[start:start + block_size, :] = R_rot @ block
    return Lm_rotated_padded[:m, :]

# ================================================
# 6. SVD DISTILLATION
# ================================================
print("\n" + "=" * 60)
print(f"6. SVD Block Rotation Distillation (θ={np.degrees(THETA):.1f}°)...")
print("=" * 60)

def svd_distillation(X, y,
                     m=NUM_PER_CLASS,
                     energy_threshold=SVD_ENERGY_THRESHOLD,
                     theta=THETA,
                     block_size=BLOCK_SIZE):
    classes       = np.unique(y)
    final_X       = []
    final_y       = []
    energy_curves = {}

    print(f"\n  θ = {np.degrees(theta):.1f}°  |  block={block_size}x{block_size}  "
          f"|  m={m}/class  |  energy={energy_threshold*100:.1f}%")
    print(f"\n{'Class':<6} {'Name':<14} {'Input':<8} {'Energy%':<10} "
          f"{'Rank':<8} {'Selected'}")
    print("-" * 58)

    for cls in classes:
        mask  = (y == cls)
        X_cls = X[mask].astype(np.float32)
        n_in  = X_cls.shape[0]

        U, S, Vt = np.linalg.svd(X_cls, full_matrices=False)

        cumulative      = np.cumsum(S ** 2) / (S ** 2).sum()
        rank            = int(np.searchsorted(cumulative, energy_threshold)) + 1
        rank            = min(rank, len(S))
        energy_captured = cumulative[rank - 1]
        energy_curves[cls] = cumulative

        n_select    = min(m, n_in)
        Lm          = U[:n_select, :rank]
        Sigma_m     = S[:rank]
        Vt_m        = Vt[:rank, :]

        Lm_rotated  = block_rotate(Lm, theta=theta, block_size=block_size)
        X_distilled = (Lm_rotated * Sigma_m) @ Vt_m

        final_X.append(X_distilled)
        final_y.append(np.full(n_select, cls))

        print(f"{cls:<6} {CIFAR_CLASSES[cls]:<14} {n_in:<8} "
              f"{energy_captured*100:<10.2f} {rank:<8} {n_select}")

    final_X = np.vstack(final_X)
    final_y = np.concatenate(final_y)
    print(f"\nDistillation complete → {final_X.shape[0]} distilled images")
    print(f"Class distribution: {np.bincount(final_y)}")

    plt.figure(figsize=(10, 5))
    for cls, curve in energy_curves.items():
        plt.plot(range(1, len(curve) + 1), curve,
                 label=CIFAR_CLASSES[cls], linewidth=1.2)
    plt.axhline(y=energy_threshold, color='black', linestyle='--',
                linewidth=1.5, label=f'{energy_threshold*100:.1f}% threshold')
    plt.xlabel('Rank'); plt.ylabel('Cumulative Energy')
    plt.title(f'SVD Energy per Class (θ={np.degrees(theta):.1f}°)')
    plt.legend(loc='lower right', fontsize=8)
    plt.tight_layout()
    plt.savefig('svd_energy_plot.png', dpi=150)
    plt.show()

    return final_X, final_y


X_distilled, y_distilled = svd_distillation(
    X_pca_distilled, y_pca_distilled,
    m=NUM_PER_CLASS,
    energy_threshold=SVD_ENERGY_THRESHOLD,
    theta=THETA,
    block_size=BLOCK_SIZE
)

# ================================================
# 7. PREPARE CNN INPUTS
# ================================================
print("\n" + "=" * 60)
print(f"7. Preparing CNN inputs (upsample to {IMG_SIZE}x{IMG_SIZE})...")
print("=" * 60)

def rescale_to_01(X_flat):
    X     = X_flat.copy()
    mins  = X.min(axis=1, keepdims=True)
    maxs  = X.max(axis=1, keepdims=True)
    denom = np.where((maxs - mins) == 0, 1e-8, maxs - mins)
    return (X - mins) / denom

def upsample(X_img, size=IMG_SIZE):
    return tf.image.resize(
        tf.constant(X_img, dtype=tf.float32),
        [size, size], method='bilinear'
    ).numpy()

X_dist_01   = rescale_to_01(X_distilled).reshape(-1, 32, 32, 3).astype(np.float32)
X_train_cnn = upsample(X_dist_01)
y_train_cnn = y_distilled.astype(np.int64)

X_val_cnn  = upsample(X_val_img)
X_test_cnn = upsample(X_test_img)
y_val_cnn  = y_val.astype(np.int64)
y_test_cnn = y_test.astype(np.int64)

y_train_cat = to_categorical(y_train_cnn, 10)
y_val_cat   = to_categorical(y_val_cnn,   10)

print(f"Distilled train : {X_train_cnn.shape}  <- {NUM_PER_CLASS*10} distilled images")
print(f"Validation      : {X_val_cnn.shape}    <- {X_val_img.shape[0]} original images")
print(f"Test            : {X_test_cnn.shape}   <- {X_test_img.shape[0]} original images")

# ================================================
# 8. tf.data PIPELINE WITH CUTOUT
# ================================================
print("\n" + "=" * 60)
print("8. Building tf.data pipeline with CutOut...")
print("=" * 60)

augment_layer = tf.keras.Sequential([
    layers.RandomFlip("horizontal"),
    layers.RandomRotation(0.055),
    layers.RandomTranslation(0.15, 0.15),
    layers.RandomZoom(0.15),
], name="augmentation")


@tf.function
def cutout_tf(image, pad_size=8):
    apply = tf.random.uniform(()) > 0.5
    if not apply:
        return image
    h  = tf.shape(image)[0]
    w  = tf.shape(image)[1]
    cy = tf.random.uniform((), 0, h, dtype=tf.int32)
    cx = tf.random.uniform((), 0, w, dtype=tf.int32)
    y1 = tf.maximum(0, cy - pad_size)
    y2 = tf.minimum(h, cy + pad_size)
    x1 = tf.maximum(0, cx - pad_size)
    x2 = tf.minimum(w, cx + pad_size)
    top    = tf.ones([y1, w, 1], dtype=tf.float32)
    mid    = tf.concat([
        tf.ones([y2 - y1, x1, 1]),
        tf.zeros([y2 - y1, x2 - x1, 1]),
        tf.ones([y2 - y1, w - x2, 1])
    ], axis=1)
    bottom = tf.ones([h - y2, w, 1], dtype=tf.float32)
    mask   = tf.concat([top, mid, bottom], axis=0)
    return image * mask


@tf.function
def preprocess_train(image, label):
    image = augment_layer(image, training=True)
    image = cutout_tf(image, pad_size=8)
    return image, label


@tf.function
def preprocess_val(image, label):
    return image, label


def make_train_dataset(X, y, batch_size=BATCH_SIZE):
    ds = tf.data.Dataset.from_tensor_slices((X, y))
    ds = ds.shuffle(len(X), reshuffle_each_iteration=True)
    ds = ds.map(preprocess_train, num_parallel_calls=tf.data.AUTOTUNE)
    ds = ds.batch(batch_size)
    ds = ds.repeat()
    ds = ds.prefetch(tf.data.AUTOTUNE)
    return ds


def make_val_dataset(X, y, batch_size=128):
    ds = tf.data.Dataset.from_tensor_slices((X, y))
    ds = ds.map(preprocess_val, num_parallel_calls=tf.data.AUTOTUNE)
    ds = ds.batch(batch_size)
    ds = ds.prefetch(tf.data.AUTOTUNE)
    return ds


train_ds = make_train_dataset(X_train_cnn, y_train_cat, batch_size=BATCH_SIZE)
val_ds   = make_val_dataset(X_val_cnn, y_val_cat)

steps_per_epoch = max(1, len(X_train_cnn) // BATCH_SIZE)
print(f"  Batch size      : {BATCH_SIZE}")
print(f"  Steps per epoch : {steps_per_epoch}")
print(f"  CutOut          : 16x16 patch, 50% probability")

# ================================================
# 9. BUILD MODEL
# ================================================
print("\n" + "=" * 60)
print("9. Building ResNet50V2 model...")
print("=" * 60)

def build_resnet_model(input_shape=(IMG_SIZE, IMG_SIZE, 3), num_classes=10):
    inputs = layers.Input(shape=input_shape)
    x      = tf.keras.applications.resnet_v2.preprocess_input(inputs * 255.0)
    base   = ResNet50V2(input_shape=input_shape,
                        include_top=False, weights='imagenet')
    base.trainable = False
    x = base(x, training=False)
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dense(512, activation='relu',
                     kernel_regularizer=tf.keras.regularizers.l2(1e-3))(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dropout(0.5)(x)
    x = layers.Dense(256, activation='relu',
                     kernel_regularizer=tf.keras.regularizers.l2(1e-3))(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dropout(0.5)(x)
    outputs = layers.Dense(num_classes, activation='softmax')(x)
    return models.Model(inputs, outputs), base


model, base_model = build_resnet_model()
model.compile(
    optimizer=tf.keras.optimizers.Adam(1e-3),
    loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.1),
    metrics=['accuracy']
)
model.summary()
trainable = sum([tf.size(w).numpy() for w in model.trainable_weights])
print(f"\nTrainable params : {trainable:,}  (head only — base frozen)")

# ================================================
# 10. TRAIN
# ================================================
print("\n" + "=" * 60)
print(f"10. Training on {NUM_PER_CLASS*10} distilled images ({EPOCHS} epochs max)...")
print("=" * 60)

cb_list = [
    callbacks.ModelCheckpoint(
        'best_resnet_distilled.keras',
        monitor='val_accuracy',
        save_best_only=True,
        verbose=1
    ),
    callbacks.EarlyStopping(
        monitor='val_accuracy',
        patience=30,
        restore_best_weights=True,
        verbose=1
    ),
    callbacks.ReduceLROnPlateau(
        monitor='val_loss',
        factor=0.5,
        patience=6,
        min_lr=1e-6,
        verbose=1
    )
]

history = model.fit(
    train_ds,
    steps_per_epoch=steps_per_epoch,
    validation_data=val_ds,
    epochs=EPOCHS,
    verbose=1
)

print("\nLoading best weights...")
model.load_weights('best_resnet_distilled.keras')
best_val_acc = max(history.history['val_accuracy'])
print(f"Best val_accuracy : {best_val_acc*100:.2f}%")

# ================================================
# 11. MEMORY-EFFICIENT TTA EVALUATION
#     FIX: processes one batch at a time
#     avoids MemoryError on 10,000 test images
# ================================================
print("\n" + "=" * 60)
print(f"11. Evaluating with TTA ({TTA_STEPS} passes)...")
print("=" * 60)

tta_augment = tf.keras.Sequential([
    layers.RandomFlip("horizontal"),
    layers.RandomRotation(0.04),
    layers.RandomTranslation(0.10, 0.10),
    layers.RandomZoom(0.10),
], name="tta_augmentation")


def tta_predict_fast(model, X, n_augments=TTA_STEPS, batch_size=128):
    """
    Memory-efficient TTA — processes one batch at a time.
    Fixes MemoryError when test set is 10,000 images at 128x128.
    """
    n     = len(X)
    preds = np.zeros((n, 10), dtype=np.float32)

    # Pass 1 — original images
    for start in range(0, n, batch_size):
        end     = min(start + batch_size, n)
        batch_x = tf.constant(X[start:end], dtype=tf.float32)
        preds[start:end] += model(batch_x, training=False).numpy()
    print(f"  Pass 1/{n_augments} (original) done...")

    # Passes 2 to N — augmented
    for step in range(n_augments - 1):
        for start in range(0, n, batch_size):
            end     = min(start + batch_size, n)
            batch_x = tf.constant(X[start:end], dtype=tf.float32)
            batch_x = tta_augment(batch_x, training=True)
            preds[start:end] += model(batch_x, training=False).numpy()
        if (step + 2) % 5 == 0 or (step + 2) == n_augments:
            print(f"  Pass {step+2}/{n_augments} done...")

    return (preds / n_augments).argmax(axis=1)


# Free memory before evaluation
gc.collect()
tf.keras.backend.clear_session()
model.load_weights('best_resnet_distilled.keras')

print(f"Running TTA on {len(X_test_cnn)} test images...")
y_pred_tta = tta_predict_fast(model, X_test_cnn,
                               n_augments=TTA_STEPS, batch_size=128)
y_pred_std = model.predict(X_test_cnn, verbose=0, batch_size=128).argmax(axis=1)

acc_std = accuracy_score(y_test_cnn, y_pred_std)
acc_tta = accuracy_score(y_test_cnn, y_pred_tta)

# ================================================
# 12. FINAL RESULTS
# ================================================
print("\n" + "=" * 60)
print("12. Final Results...")
print("=" * 60)

print(f"\nRotation angle                 : θ = {np.degrees(THETA):.1f}°")
print(f"SVD energy threshold           : {SVD_ENERGY_THRESHOLD*100:.1f}%")
print(f"Best val_accuracy (training)   : {best_val_acc*100:.2f}%")
print(f"\nStandard prediction (1 pass)   : {acc_std * 100:.2f}%")
print(f"TTA prediction    ({TTA_STEPS} passes): {acc_tta * 100:.2f}%")
print(f"TTA gain                       : +{(acc_tta - acc_std) * 100:.2f}%")

print(f"\n{'=' * 60}")
print(f"FINAL TEST ACCURACY (TTA x{TTA_STEPS}) : {acc_tta * 100:.2f}%")
print(f"Model               : ResNet50V2 (ImageNet, frozen head)")
print(f"Distillation        : PCA ({VARIANCE_THRESHOLD*100:.1f}% var) + SVD Block Rotation "
      f"({SVD_ENERGY_THRESHOLD*100:.1f}% energy, θ={np.degrees(THETA):.1f}°)")
print(f"Trained on          : {X_train_cnn.shape[0]} distilled images")
print(f"Tested on           : {X_test_cnn.shape[0]} original images "
      f"({X_test_cnn.shape[0]//10}/class)")
print(f"{'=' * 60}")

print(f"\nClassification Report (TTA x{TTA_STEPS}):")
print(classification_report(
    y_test_cnn, y_pred_tta,
    target_names=CIFAR_CLASSES,
    digits=4
))

# ================================================
# 13. CONFUSION MATRIX
# ================================================
print("\nPlotting confusion matrix...")

cm = confusion_matrix(y_test_cnn, y_pred_tta)

plt.figure(figsize=(12, 9))
sns.heatmap(
    cm,
    annot=True,
    fmt='d',
    cmap='Blues',
    xticklabels=CIFAR_CLASSES,
    yticklabels=CIFAR_CLASSES,
    linewidths=0.5,
    annot_kws={"size": 9}
)
plt.title(
    f'Confusion Matrix — CIFAR-10 — TTA x{TTA_STEPS} ({acc_tta*100:.2f}%)\n'
    f'θ={np.degrees(THETA):.0f}°  |  500 distilled images  |  10,000 test images',
    fontsize=12, pad=15
)
plt.ylabel('True Label', fontsize=11)
plt.xlabel('Predicted Label', fontsize=11)
plt.xticks(rotation=45, ha='right')
plt.yticks(rotation=0)
plt.tight_layout()
plt.savefig('confusion_matrix_cifar10.png', dpi=150, bbox_inches='tight')
plt.show()
print("Confusion matrix saved as 'confusion_matrix_cifar10.png'")
print("\nAll done.")

# ================================================
# 14. TRAINING HISTORY — ACCURACY & LOSS CURVES
# ================================================
print("\nPlotting training history...")

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# ── Accuracy ──
axes[0].plot(history.history['accuracy'],
             label='Train Accuracy', color='#3B6EB5', linewidth=1.8)
axes[0].plot(history.history['val_accuracy'],
             label='Val Accuracy', color='#C0401A', linewidth=1.8, linestyle='--')
axes[0].set_title('Model Accuracy', fontsize=13, fontweight='bold')
axes[0].set_xlabel('Epoch', fontsize=11)
axes[0].set_ylabel('Accuracy', fontsize=11)
axes[0].legend(fontsize=10)
axes[0].grid(True, alpha=0.3)

# ── Loss ──
axes[1].plot(history.history['loss'],
             label='Train Loss', color='#3B6EB5', linewidth=1.8)
axes[1].plot(history.history['val_loss'],
             label='Val Loss', color='#C0401A', linewidth=1.8, linestyle='--')
axes[1].set_title('Model Loss', fontsize=13, fontweight='bold')
axes[1].set_xlabel('Epoch', fontsize=11)
axes[1].set_ylabel('Loss', fontsize=11)
axes[1].legend(fontsize=10)
axes[1].grid(True, alpha=0.3)

fig.suptitle(
    f'Training History — CIFAR-10 — TTA x{TTA_STEPS} ({acc_tta*100:.2f}%)\n'
    f'θ={np.degrees(THETA):.0f}°  |  {NUM_PER_CLASS*10} distilled images',
    fontsize=12, fontweight='bold'
)
plt.tight_layout()
plt.savefig('training_history_cifar10.png', dpi=150, bbox_inches='tight')
plt.show()
print("Training history saved as 'training_history_cifar10.png'")
print("\nAll done.")