import torch
import torch.nn as nn
from src.models.student import LIPEV2StudentGaze360Gold, LIPEV2StudentFinal

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def calc_flops_final(model, input_size):
    total_flops = 0
    batch, patches, h, w = input_size # (1, 4, 16, 16)
    
    # 1. Appearance Branch (LIPEFinalAppearance)
    # Conv1: 1 -> 32, 3x3, p=1
    total_flops += patches * (2 * (3 * 3 * 1) * 32 * 16 * 16)
    # Conv2: 32 -> 64, 3x3, s=2, p=1
    total_flops += patches * (2 * (3 * 3 * 32) * 64 * 8 * 8)
    # Conv3: 64 -> 128, 3x3, s=2, p=1
    total_flops += patches * (2 * (3 * 3 * 64) * 128 * 4 * 4)
    # Conv4: 128 -> 256, 3x3, p=1
    total_flops += patches * (2 * (3 * 3 * 128) * 256 * 4 * 4)
    
    # 2. Pooling & Shape Invariance
    # GAP/GMP: (256 * 4 * 4) * 2 (avg/max)
    total_flops += patches * (2 * 256 * 4 * 4)
    # Linear Projection: 512 -> 512
    total_flops += patches * (2 * 512 * 512)
    
    # 3. Fusion & Regression
    # Fusion: (512 + 128) -> 512
    total_flops += 2 * (512 + 128) * 512
    # Regression: 512 -> 2
    total_flops += 2 * 512 * 2
    
    return total_flops

# Current GOLD
model_gold = LIPEV2StudentGaze360Gold()
# Final Architecture
model_final = LIPEV2StudentFinal()

print(f"--- ARCHITECTURAL SYNC REPORT ---")
print(f"Model: LIPEV2StudentFinal")
print(f"Total Parameters: {count_parameters(model_final):,}")
flops_final = calc_flops_final(model_final, (1, 4, 16, 16))
print(f"FLOPs (16x16 End-to-End): {flops_final/1e6:.2f} MFLOPs")
print(f"FLOPs in GFLOPs: {flops_final/1e9:.4f} GFLOPs")

print(f"\n--- FOR REFERENCE: OLD GOLD STATS ---")
print(f"Total Parameters: {count_parameters(model_gold):,}")