"""Predict next frame using improved U-Net v2 (fp16 weights)."""
import json
import sys
from typing import Any

import numpy as np
import torch

sys.path.insert(0, '/home/coder/code')
from unet_v3 import UNetV3


def load_model(model_dir: str) -> Any:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    with open(f"{model_dir}/config.json", 'r') as f:
        config = json.load(f)

    model = UNetV3(
        in_channels=config['in_channels'],
        out_channels=config['out_channels'],
        enc_channels=config['enc_channels'],
        dec_channels=config['dec_channels'],
        use_tanh=config.get('use_tanh', False),
    )

    state_fp16 = torch.load(f"{model_dir}/model.pt", map_location='cpu', weights_only=True)
    state_fp32 = {k: v.float() for k, v in state_fp16.items()}
    model.load_state_dict(state_fp32)
    model.to(device)
    model.eval()

    return {'model': model, 'device': device, 'n_context': config['n_context']}


def predict_next_frame(model_dict: Any, context_frames: np.ndarray) -> np.ndarray:
    model = model_dict['model']
    device = model_dict['device']
    n_ctx = model_dict['n_context']

    # Use last n_ctx frames, pad if needed
    if len(context_frames) >= n_ctx:
        frames = context_frames[-n_ctx:]
    else:
        pad_count = n_ctx - len(context_frames)
        padding = np.stack([context_frames[0]] * pad_count, axis=0)
        frames = np.concatenate([padding, context_frames], axis=0)

    # Normalize to [0, 1], convert to CHW
    frames_norm = frames.astype(np.float32) / 255.0
    last_frame = frames_norm[-1]

    # Stack context: (n_ctx, 3, H, W) -> (n_ctx*3, H, W)
    stacked = np.transpose(frames_norm, (0, 3, 1, 2)).reshape(-1, 64, 64)
    last_chw = np.transpose(last_frame, (2, 0, 1))

    with torch.no_grad():
        inp = torch.from_numpy(stacked).unsqueeze(0).to(device)
        last_t = torch.from_numpy(last_chw).unsqueeze(0).to(device)
        delta = model(inp)
        pred = torch.clamp(last_t + delta, 0, 1)
        pred = pred.squeeze(0).cpu().numpy()

    # Convert to uint8 HWC
    pred = np.clip(pred, 0, 1)
    pred = (pred * 255.0).astype(np.uint8)
    return np.transpose(pred, (1, 2, 0))