#!/usr/bin/env python3
# Evaluate a trained SeqGRU on the validation set; reads input_dim from meta.json

import os, json, argparse
import numpy as np
import torch, torch.nn as nn
from sklearn.metrics import classification_report, confusion_matrix

class SeqGRU(nn.Module):
    def __init__(self, input_dim, hidden=128, num_classes=26):
        super().__init__()
        self.gru = nn.GRU(input_dim, hidden, batch_first=True, bidirectional=True)
        self.head = nn.Sequential(
            nn.Linear(hidden*2, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, num_classes),
        )
    def forward(self, x):
        h,_ = self.gru(x)
        return self.head(h[:, -1, :])

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--landmarks", default="landmarks_seq32")
    ap.add_argument("--model", required=True)
    args = ap.parse_args()

    vaX = np.load(os.path.join(args.landmarks,"val_X.npy"))
    vaY = np.load(os.path.join(args.landmarks,"val_y.npy"))
    classes = json.load(open(os.path.join(args.landmarks,"class_names.json")))
    meta = json.load(open(os.path.join(args.landmarks,"meta.json")))
    T = int(meta.get("frames", vaX.shape[1]))
    input_dim = int(meta.get("input_dim", vaX.shape[-1]))

    state = torch.load(args.model, map_location="cpu", weights_only=False)
    X_mean, X_std = state["X_mean"], state["X_std"]
    if isinstance(X_mean, torch.Tensor): X_mean = X_mean.numpy()
    if isinstance(X_std,  torch.Tensor): X_std  = X_std.numpy()
    X_mean = X_mean.astype(np.float32)
    X_std  = (X_std.astype(np.float32) + 1e-6)

    vaXn = (vaX - X_mean) / X_std

    device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
    model = SeqGRU(input_dim=input_dim, hidden=128, num_classes=len(classes))
    model.load_state_dict(state["model"])
    model.eval().to(device)

    with torch.no_grad():
        xb = torch.from_numpy(vaXn).float().to(device)
        logits = model(xb)
        pred = logits.argmax(1).cpu().numpy()

    cm = confusion_matrix(vaY, pred)
    print("Classes:", classes)
    print("\nConfusion matrix (rows=true, cols=pred):\n", cm)
    print("\nReport:\n", classification_report(vaY, pred, target_names=classes))

if __name__ == "__main__":
    main()