"""Synth6 = Synth4 (EarthView negatives, p_ev) + HOT building scenes (p_hot of building scenes) + Korean-style vinyl greenhouse synthesis (p_gh)."""
import os, time, math
import cv2, numpy as np, torch
from train2 import CROP
from train3 import degrade
from ev_synth import Synth4

class Synth6(Synth4):
    def __init__(self, n, p_ev=0.2, p_hot=0.5, p_gh=0.12):
        super().__init__(n, p_ev)
        self.p_hot, self.p_gh = p_hot, p_gh
        self.hx = np.load("/workspace/prep/hot_x.npy", mmap_mode="r"); self.hy = np.load("/workspace/prep/hot_y.npy", mmap_mode="r")
        print("HOT scenes", self.hx.shape, flush=True)

    def scene(self, rng, kind):
        if kind == "b" and rng.random() < self.p_hot:
            i = rng.integers(len(self.hx)); s = rng.uniform(0.85, 1.15); size = min(256, int(round(CROP / s)))
            r, c = rng.integers(0, 257 - size, 2)
            x = cv2.resize(np.ascontiguousarray(self.hx[i, r:r + size, c:c + size]), (CROP, CROP), interpolation=cv2.INTER_AREA)
            y = cv2.resize(np.ascontiguousarray(self.hy[i, r:r + size, c:c + size]), (CROP, CROP), interpolation=cv2.INTER_NEAREST)
            return x.copy(), y.astype(np.uint8)          # code 1 = building (others unlabeled)
        return super().scene(rng, kind)

    @staticmethod
    def greenhouses(rng, shape):
        """Rows of long, narrow tunnel greenhouses: returns (mask, rendered rgb layer)."""
        m = np.zeros(shape, np.uint8); lay = np.zeros(shape + (3,), np.float32)
        th = rng.uniform(0, np.pi); u = np.array([np.cos(th), np.sin(th)]); v = np.array([-u[1], u[0]])
        n = int(rng.integers(2, 9)); w = rng.uniform(8, 14); gap = rng.uniform(1.5, 5); L = rng.uniform(30, 140)
        c0 = np.array([rng.uniform(40, shape[1] - 40), rng.uniform(40, shape[0] - 40)])
        base = rng.uniform(185, 245); tint = rng.uniform(-12, 12, 3)
        yy, xx = np.mgrid[:shape[0], :shape[1]].astype(np.float32)
        for k in range(n):
            ck = c0 + v * (k - n / 2) * (w + gap) + u * rng.uniform(-6, 6)
            Lk = L * rng.uniform(0.8, 1.0)
            pts = np.array([ck + u * Lk / 2 + v * w / 2, ck + u * Lk / 2 - v * w / 2, ck - u * Lk / 2 - v * w / 2, ck - u * Lk / 2 + v * w / 2], np.int32)
            mk = np.zeros(shape, np.uint8); cv2.fillPoly(mk, [pts], 1)
            dv = ((xx - ck[0]) * v[0] + (yy - ck[1]) * v[1]) / (w / 2)          # -1..1 across the tunnel
            shade = 0.78 + 0.22 * np.cos(np.clip(dv, -1, 1) * np.pi / 2)        # bright ridge, darker sides
            for ch in range(3):
                lay[..., ch] = np.where(mk > 0, (base + tint[ch]) * shade + rng.normal(0, 3, shape), lay[..., ch])
            m |= mk
        return m, np.clip(lay, 0, 255)

    def __getitem__(self, idx):
        rng = np.random.default_rng((idx * 7919 + os.getpid() * 104729 + time.time_ns()) % 2**32)
        if rng.random() >= self.p_gh:
            return super().__getitem__(idx)
        x, y = self.scene(rng, "g" if rng.random() < 0.6 else "any")
        pre = x.copy(); sem = np.zeros(y.shape, np.uint8); sem[y == 1] = 1; sem[y == 2] = 2
        if rng.random() < 0.3:   # pre-existing greenhouses in both images (not a change)
            m0, l0 = self.greenhouses(rng, y.shape); a = (m0[..., None] * rng.uniform(0.85, 0.95))
            pre = (pre * (1 - a) + l0 * a).astype(np.uint8); sem[m0 > 0] = 1
        post = pre.copy(); m, lay = self.greenhouses(rng, y.shape)
        m[sem == 1] = 0
        a = cv2.GaussianBlur(m.astype(np.float32), (0, 0), 0.6)[..., None] * rng.uniform(0.85, 0.97)
        post = (post * (1 - a) + lay * a).astype(np.uint8)
        lb = m.copy(); lt = np.zeros_like(m); sem_pre = sem.copy(); sem_post = sem.copy(); sem_post[m > 0] = 1
        if rng.random() < 0.15:   # demolition: reversed pair is not a target
            pre, post = post, pre; sem_pre, sem_post = sem_post, sem_pre; lb[:] = 0
        pre = degrade(self.photo(pre, rng), rng); post = degrade(self.photo(post, 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)
            sem_pre = cv2.warpAffine(sem_pre, M, (CROP, CROP), flags=cv2.INTER_NEAREST, 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, lb, lt, sem_pre, sem_post = map(d4, (pre, post, lb, lt, sem_pre, sem_post))
        t = lambda a: torch.from_numpy(a.transpose(2, 0, 1).copy())
        pres = np.array([lb.sum() >= 20, 0], np.float32)
        return (t(pre), t(post), torch.from_numpy(np.stack([lb, lt]).astype(np.float32)), torch.from_numpy(pres),
                torch.from_numpy(sem_pre.astype(np.int64)), torch.from_numpy(sem_post.astype(np.int64)))