"""Real-ESRGAN video upscaler (no basicsr dependency).  usage: upscale.py in.mp4 out.mp4 [model] [out_scale]
model: general (realesr-general-x4v3, SRVGG compact) | x4plus (RRDBNet).  Output = 4x net, resized to out_scale x.
"""
import sys, os, subprocess, numpy as np, cv2, torch, torch.nn as nn, torch.nn.functional as F

W = os.environ.get("ESRGAN_WEIGHTS", os.path.join(os.path.dirname(os.path.abspath(__file__)), "weights"))


class SRVGG(nn.Module):
    def __init__(self, nf=64, nc=32, up=4):
        super().__init__()
        self.up = up
        layers = [nn.Conv2d(3, nf, 3, 1, 1), nn.PReLU(nf)]
        for _ in range(nc):
            layers += [nn.Conv2d(nf, nf, 3, 1, 1), nn.PReLU(nf)]
        layers += [nn.Conv2d(nf, 3 * up * up, 3, 1, 1)]
        self.body = nn.ModuleList(layers)
        self.upsampler = nn.PixelShuffle(up)

    def forward(self, x):
        o = x
        for l in self.body:
            o = l(o)
        return self.upsampler(o) + F.interpolate(x, scale_factor=self.up, mode="nearest")


class RDB(nn.Module):
    def __init__(self, nf=64, gc=32):
        super().__init__()
        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1); self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1)
        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1); self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1)
        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1); self.lrelu = nn.LeakyReLU(0.2, True)

    def forward(self, x):
        x1 = self.lrelu(self.conv1(x)); x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1))); x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
        return self.conv5(torch.cat((x, x1, x2, x3, x4), 1)) * 0.2 + x


class RRDB(nn.Module):
    def __init__(self, nf, gc=32):
        super().__init__(); self.rdb1 = RDB(nf, gc); self.rdb2 = RDB(nf, gc); self.rdb3 = RDB(nf, gc)

    def forward(self, x):
        return self.rdb3(self.rdb2(self.rdb1(x))) * 0.2 + x


class RRDBNet(nn.Module):
    def __init__(self, nf=64, nb=23, gc=32):
        super().__init__()
        self.conv_first = nn.Conv2d(3, nf, 3, 1, 1)
        self.body = nn.Sequential(*[RRDB(nf, gc) for _ in range(nb)])
        self.conv_body = nn.Conv2d(nf, nf, 3, 1, 1)
        self.conv_up1 = nn.Conv2d(nf, nf, 3, 1, 1); self.conv_up2 = nn.Conv2d(nf, nf, 3, 1, 1)
        self.conv_hr = nn.Conv2d(nf, nf, 3, 1, 1); self.conv_last = nn.Conv2d(nf, 3, 3, 1, 1)
        self.lrelu = nn.LeakyReLU(0.2, True)

    def forward(self, x):
        fea = self.conv_first(x)
        fea = fea + self.conv_body(self.body(fea))
        fea = self.lrelu(self.conv_up1(F.interpolate(fea, scale_factor=2, mode="nearest")))
        fea = self.lrelu(self.conv_up2(F.interpolate(fea, scale_factor=2, mode="nearest")))
        return self.conv_last(self.lrelu(self.conv_hr(fea)))


def load(model):
    dev = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    if model == "general":
        net = SRVGG(); sd = torch.load(os.path.join(W, "realesr-general-x4v3.pth"), map_location="cpu")
        sd = sd.get("params", sd)
    else:
        net = RRDBNet(); sd = torch.load(os.path.join(W, "RealESRGAN_x4plus.pth"), map_location="cpu")
        sd = sd.get("params_ema", sd)
    net.load_state_dict(sd, strict=True); net.eval().to(dev)
    return net, dev


@torch.no_grad()
def up_frame(net, dev, bgr, out_scale):
    h, w = bgr.shape[:2]
    x = torch.from_numpy(bgr[..., ::-1].copy()).permute(2, 0, 1).float().div(255).unsqueeze(0).to(dev)
    y = net(x).clamp(0, 1)[0].permute(1, 2, 0).cpu().numpy()[..., ::-1]
    y = (y * 255).round().astype(np.uint8)
    if out_scale != 4:
        y = cv2.resize(y, (int(w * out_scale), int(h * out_scale)), interpolation=cv2.INTER_AREA)
    return y


def main():
    src, dst = sys.argv[1], sys.argv[2]
    model = sys.argv[3] if len(sys.argv) > 3 else "general"
    out_scale = float(sys.argv[4]) if len(sys.argv) > 4 else 2
    frame_only = os.environ.get("FRAME")
    p = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width,height,r_frame_rate", "-of", "csv=p=0", src], capture_output=True, text=True).stdout.strip().split(",")
    w, h, fps = int(p[0]), int(p[1]), p[2]
    net, dev = load(model)
    dec = subprocess.Popen(["ffmpeg", "-v", "error", "-i", src, "-f", "rawvideo", "-pix_fmt", "bgr24", "-"], stdout=subprocess.PIPE)
    enc = None
    i = 0
    while True:
        buf = dec.stdout.read(w * h * 3)
        if len(buf) < w * h * 3: break
        if frame_only is not None and i != int(frame_only):
            i += 1; continue
        f = np.frombuffer(buf, np.uint8).reshape(h, w, 3)
        y = up_frame(net, dev, f, out_scale)
        if frame_only is not None:
            cv2.imwrite(dst, y); return
        if enc is None:
            enc = subprocess.Popen(["ffmpeg", "-v", "error", "-y", "-f", "rawvideo", "-pix_fmt", "bgr24", "-s", f"{y.shape[1]}x{y.shape[0]}", "-r", fps, "-i", "-",
                                    "-c:v", "libx264", "-crf", "12", "-preset", "slow", "-pix_fmt", "yuv420p", dst], stdin=subprocess.PIPE)
        enc.stdin.write(y.tobytes()); i += 1
        if i % 24 == 0: print(f"{os.path.basename(src)} {i}", flush=True)
    enc.stdin.close(); enc.wait(); print(os.path.basename(src), "done", i, "frames")


if __name__ == "__main__":
    main()