"""SpeedNet: optical-flow -> speed regression for onboard/POV footage.

Input per timestep: Farneback optical flow between consecutive frames of
320x180 @ 15 fps video, resized to 128x72, 2 channels (px/frame at 320x180
scale), divided by FLOW_SCALE. Output: speed in m/s after * SPEED_SCALE.

IMPORTANT — inference protocol:
Run the GRU in STREAMING mode (carry hidden state across chunks, see
`predict_streaming`). Resetting the hidden state every few frames collapses
indoor-domain predictions to ~half scale; the domain-scale calibration lives
in long GRU context.
"""
import numpy as np
import torch
import torch.nn as nn

FLOW_SCALE = 8.0
SPEED_SCALE = 40.0


class SpeedNet(nn.Module):
    def __init__(self, hidden=192):
        super().__init__()
        def blk(i, o):
            return [nn.Conv2d(i, o, 3, 2, 1), nn.BatchNorm2d(o), nn.SiLU()]
        self.cnn = nn.Sequential(*blk(2, 32), *blk(32, 64), *blk(64, 128),
                                 *blk(128, 256), nn.AdaptiveAvgPool2d(1),
                                 nn.Flatten())
        self.gru = nn.GRU(256, hidden, num_layers=2, batch_first=True)
        self.head = nn.Linear(hidden, 1)

    def forward(self, x):  # x: [B, T, 2, 72, 128], normalized by FLOW_SCALE
        B, T = x.shape[:2]
        f = self.cnn(x.flatten(0, 1)).view(B, T, -1)
        h, _ = self.gru(f)
        return self.head(h).squeeze(-1)  # [B, T], * SPEED_SCALE -> m/s


@torch.no_grad()
def predict_streaming(model, flow, device='cuda', chunk=240):
    """flow: [N, 72, 128, 2] float (px/frame). Returns speed m/s per frame."""
    model.eval()
    h = None
    out = []
    for s in range(0, len(flow), chunk):
        x = torch.from_numpy(flow[s:s + chunk].astype(np.float32)
                             .transpose(0, 3, 1, 2)).to(device) / FLOW_SCALE
        f = model.cnn(x)[None]
        y, h = model.gru(f, h)
        out.append(model.head(y)[0, :, 0].cpu().numpy() * SPEED_SCALE)
    return np.concatenate(out)


def compute_flow(video_path):
    """Farneback flow matching the training distribution.
    Video should be (or be resized to) 320x180 @ 15 fps."""
    import cv2
    cap = cv2.VideoCapture(video_path)
    prev, flows = None, []
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        if frame.shape[:2] != (180, 320):
            frame = cv2.resize(frame, (320, 180))
        g = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        if prev is not None:
            fl = cv2.calcOpticalFlowFarneback(prev, g, None,
                                              0.5, 3, 15, 3, 5, 1.2, 0)
            flows.append(cv2.resize(fl, (128, 72)))
        prev = g
    cap.release()
    return np.stack(flows)