"""Synth4: Synth3 synthetic pairs + real EarthView single-date 'no change' pairs (fraction P_EV)."""
import os
import time

import cv2
import numpy as np
import torch

from train2 import CROP
from train3 import Synth3, degrade


class Synth4(Synth3):
    def __init__(self, n, p_ev=0.2):
        super().__init__(n)
        self.p_ev = p_ev
        self.ev = np.load("/workspace/prep/ev_imgs.npy", mmap_mode="r")
        print("earthview singles", self.ev.shape, flush=True)

    def __getitem__(self, idx):
        rng = np.random.default_rng((idx * 7919 + os.getpid() * 104729 + time.time_ns()) % 2**32)
        if rng.random() >= self.p_ev:
            return super().__getitem__(idx)
        img = self.ev[rng.integers(len(self.ev))]
        size = int(CROP * rng.uniform(0.55, 0.7))
        r, c = rng.integers(0, img.shape[0] - size + 1, 2)
        crop = cv2.resize(np.ascontiguousarray(img[r:r + size, c:c + size]), (CROP, CROP), interpolation=cv2.INTER_CUBIC)
        pre, post = degrade(self.photo(crop, rng), rng), degrade(self.photo(crop.copy(), rng), rng)
        if rng.random() < 0.8:
            M = cv2.getRotationMatrix2D((CROP / 2, CROP / 2), rng.uniform(-1.5, 1.5), rng.uniform(0.98, 1.02))
            M[:, 2] += rng.uniform(-4, 4, 2)
            pre = cv2.warpAffine(pre, M, (CROP, CROP), borderMode=cv2.BORDER_REFLECT)
        k = rng.integers(8)
        def d4(a):
            a = np.rot90(a, k % 4)
            return np.ascontiguousarray(a[:, ::-1] if k >= 4 else a)
        pre, post = d4(pre), d4(post)
        ig = np.full((CROP, CROP), 255, np.int64)
        t = lambda a: torch.from_numpy(a.transpose(2, 0, 1).copy())
        return (t(pre), t(post), torch.zeros(2, CROP, CROP), torch.zeros(2), torch.from_numpy(ig), torch.from_numpy(ig.copy()))