import os
import sys
import random
import math
import torch
import numpy as np
import kornia
import nvdiffrast.torch as dr
import torch.nn.functional as F
from argparse import ArgumentParser

from core.registry import register_method
from core.base_method import BaseMethod

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GI-GS_offy')))

from utils.loss_utils import l1_loss, ssim
from utils.image_utils import psnr
from gaussian_renderer import render
from diff_gaussian_rasterization import Gaussian_SSR
from pbr import CubemapLight, get_brdf_lut, pbr_shading
from scene import Scene, GaussianModel
from arguments import ModelParams, PipelineParams, OptimizationParams

def linear_to_srgb(linear):
    eps = torch.finfo(torch.float32).eps
    srgb0 = 323 / 25 * linear
    srgb1 = (211 * torch.clamp(linear, min=eps) ** (5 / 12) - 11) / 200
    return torch.where(linear <= 0.0031308, srgb0, srgb1)

def srgb_to_linear(srgb):
    linear0 = 25 / 323 * srgb
    linear1 = ((srgb + 0.055) / 1.055)**2.4
    return torch.where(srgb <= 0.04045, linear0, linear1)

def get_tv_loss(gt_image, prediction, pad=1, step=1):
    if pad > 1:
        gt_image = F.avg_pool2d(gt_image, pad, pad)
        prediction = F.avg_pool2d(prediction, pad, pad)
    rgb_grad_h = torch.exp(-(gt_image[:, 1:, :] - gt_image[:, :-1, :]).abs().mean(dim=0, keepdim=True))
    rgb_grad_w = torch.exp(-(gt_image[:, :, 1:] - gt_image[:, :, :-1]).abs().mean(dim=0, keepdim=True))
    tv_h = torch.pow(prediction[:, 1:, :] - prediction[:, :-1, :], 2)
    tv_w = torch.pow(prediction[:, :, 1:] - prediction[:, :, :-1], 2)
    tv_loss = (tv_h * rgb_grad_h).mean() + (tv_w * rgb_grad_w).mean()
    if step > 1:
        for s in range(2, step + 1):
            rgb_grad_h = torch.exp(-(gt_image[:, s:, :] - gt_image[:, :-s, :]).abs().mean(dim=0, keepdim=True))
            rgb_grad_w = torch.exp(-(gt_image[:, :, s:] - gt_image[:, :, :-s]).abs().mean(dim=0, keepdim=True))
            tv_h = torch.pow(prediction[:, s:, :] - prediction[:, :-s, :], 2)
            tv_w = torch.pow(prediction[:, :, s:] - prediction[:, :, :-s], 2)
            tv_loss += (tv_h * rgb_grad_h).mean() + (tv_w * rgb_grad_w).mean()
    return tv_loss

def get_masked_tv_loss(mask, gt_image, prediction):
    rgb_grad_h = torch.exp(-(gt_image[:, 1:, :] - gt_image[:, :-1, :]).abs().mean(dim=0, keepdim=True))
    rgb_grad_w = torch.exp(-(gt_image[:, :, 1:] - gt_image[:, :, :-1]).abs().mean(dim=0, keepdim=True))
    tv_h = torch.pow(prediction[:, 1:, :] - prediction[:, :-1, :], 2)
    tv_w = torch.pow(prediction[:, :, 1:] - prediction[:, :, :-1], 2)
    mask = mask.float()
    mask_h = mask[:, 1:, :] * mask[:, :-1, :]
    mask_w = mask[:, :, 1:] * mask[:, :, :-1]
    tv_loss = (tv_h * rgb_grad_h * mask_h).mean() + (tv_w * rgb_grad_w * mask_w).mean()
    return tv_loss

def get_envmap_dirs(res=[512, 1024]):
    gy, gx = torch.meshgrid(
        torch.linspace(0.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device="cuda"),
        torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device="cuda"),
        indexing="ij",
    )
    sintheta, costheta = torch.sin(gy * np.pi), torch.cos(gy * np.pi)
    sinphi, cosphi = torch.sin(gx * np.pi), torch.cos(gx * np.pi)
    reflvec = torch.stack((sintheta * sinphi, costheta, -sintheta * cosphi), dim=-1)
    return reflvec

@register_method("gigs")
class GIGSWrapper(BaseMethod):
    def __init__(self, dataset_config, hyperparams):
        self.parser = ArgumentParser()
        self.lp = ModelParams(self.parser)
        self.op = OptimizationParams(self.parser)
        self.pp = PipelineParams(self.parser)
        self.args = self.parser.parse_args([])

        self.args.source_path = dataset_config["source_path"]
        self.args.model_path = dataset_config["model_path"]
        self.args.resolution = dataset_config.get("resolution", 1)
        self.args.eval = True
        self.args.metallic = True
        self.args.indirect = True
        self.args.pbr_iteration = 7000
        self.args.tone = False
        self.args.gamma = False
        self.args.normal_tv = 1.0
        self.args.brdf_tv = 1.0
        self.args.env_tv = 0.01
        self.args.radius = 0.8
        self.args.bias = 0.01
        self.args.thick = 0.05
        self.args.delta = 0.0625
        self.args.step = 16
        self.args.start = 8

        if "360" in str(self.args.source_path).lower() and "indoor" not in str(self.args.source_path).lower():
            self.args.degree = 3
        else:
            self.args.degree = 1

        self.track_decoupling = hyperparams.get("track_decoupling", False)

        self.dataset = self.lp.extract(self.args)
        self.dataset.sh_degree = self.args.degree
        self.opt = self.op.extract(self.args)
        self.pipe = self.pp.extract(self.args)

        self.gaussians = GaussianModel(self.dataset.sh_degree)
        # INJECTED_RES_FIX begin
        import sys as _sys
        _scene, _explicit_res = None, None
        for _i, _a in enumerate(_sys.argv[:-1]):
            _v = _sys.argv[_i + 1]
            if _a == "--scene":          _scene = _v
            elif _a == "--source_path":  _scene = _v.rstrip("/").split("/")[-1]
            elif _a == "--resolution":
                try: _explicit_res = int(_v)
                except: pass
        _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"}
        if _explicit_res is not None and _explicit_res > 0:
            _res = _explicit_res
        elif _scene is not None:
            _res = 4 if _scene in _OUTDOOR_360 else 2
        else:
            _res = None
        try:
            if _res is not None:
                self.dataset.resolution = _res
                print("[res-fix] scene=%s explicit=%s -> res=%s  (%s)" % (_scene, _explicit_res, _res, __file__))
        except Exception as _e:
            print("[res-fix] FAILED:", _e)
        # INJECTED_RES_FIX end
        self.scene = Scene(self.dataset, self.gaussians)
        self.gaussians.training_setup(self.opt)

        self.brdf_lut = get_brdf_lut().cuda()
        self.envmap_dirs = get_envmap_dirs()
        self.cubemap = CubemapLight(base_res=256).cuda()
        self.cubemap.train()

        param_groups = [{"name": "cubemap", "params": self.cubemap.parameters(), "lr": self.opt.opacity_lr}]
        self.light_optimizer = torch.optim.Adam(param_groups, lr=self.opt.opacity_lr)

        bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0]
        self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
        self.viewpoint_stack = self.scene.getTrainCameras().copy()
        self.last_n_gaussians = len(self.gaussians.get_xyz)
        self.canonical_rays = self.scene.get_canonical_rays()

    def train_iteration(self, step):
        self.gaussians.update_learning_rate(step)
        if step % 1000 == 0:
            self.gaussians.oneupSHdegree()

        if not self.viewpoint_stack:
            self.viewpoint_stack = self.scene.getTrainCameras().copy()
            
        viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1))

        try:
            c2w = torch.inverse(viewpoint_cam.world_view_transform.T)
        except:
            return {}, {}

        bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background
        current_bg = bg if step <= self.args.pbr_iteration else torch.zeros_like(bg)

        render_pkg = render(
            viewpoint_camera=viewpoint_cam,
            pc=self.gaussians,
            pipe=self.pipe,
            bg_color=current_bg,
            pad_normal=False,
            derive_normal=True,
            radius=self.args.radius,
            bias=self.args.bias,
            thick=self.args.thick,
            delta=self.args.delta,
            step=self.args.step,
            start=self.args.start
        )

        image = render_pkg["render"]
        viewspace_point_tensor = render_pkg["viewspace_points"]
        visibility_filter = render_pkg["visibility_filter"]
        radii = render_pkg["radii"]
        normal_map_from_depth = render_pkg["normal_map_from_depth"]
        normal_map = render_pkg["normal_map"]
        albedo_map = render_pkg["albedo_map"]
        roughness_map = render_pkg["roughness_map"]
        metallic_map = render_pkg["metallic_map"]

        rmax, rmin = 1.0, 0.04
        roughness_map = roughness_map * (rmax - rmin) + rmin

        H, W = viewpoint_cam.image_height, viewpoint_cam.image_width
        view_dirs = -((F.normalize(self.canonical_rays[:, None, :], p=2, dim=-1) * c2w[None, :3, :3]).sum(dim=-1).reshape(H, W, 3))

        alpha_mask = viewpoint_cam.gt_alpha_mask.cuda()
        gt_image = viewpoint_cam.original_image[0:3, :, :].cuda()
        gt_image = (gt_image * alpha_mask + current_bg[:, None, None] * (1.0 - alpha_mask)).clamp(0.0, 1.0)

        loss_target = torch.tensor(0.0, device="cuda")
        loss_parasitic = torch.tensor(0.0, device="cuda")

        if step <= self.args.pbr_iteration:
            Ll1 = F.l1_loss(image, gt_image)
            loss_target = (1.0 - self.opt.lambda_dssim) * Ll1
            ssim_val = ssim(image, gt_image)
            mask = render_pkg["normal_from_depth_mask"]
            normal_loss = F.l1_loss(normal_map[:, mask], normal_map_from_depth[:, mask])
            normal_tv_loss = get_tv_loss(gt_image, normal_map, pad=1, step=1)
            loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_val) + normal_loss + normal_tv_loss * self.args.normal_tv
            loss = loss_target + loss_parasitic
        else:
            occlusion = render_pkg["occlusion_map"].permute(1, 2, 0) if self.args.indirect else torch.ones_like(roughness_map).permute(1, 2, 0)
            normal_mask = render_pkg["normal_mask"]
            out_normal_view = render_pkg["out_normal_view"]
            depth_pos = render_pkg["depth_pos"]

            self.cubemap.build_mips()
            pbr_result = pbr_shading(
                light=self.cubemap,
                normals=normal_map.permute(1, 2, 0).detach(),
                view_dirs=view_dirs,
                mask=normal_mask.permute(1, 2, 0),
                albedo=albedo_map.permute(1, 2, 0),
                roughness=roughness_map.permute(1, 2, 0),
                metallic=metallic_map.permute(1, 2, 0) if self.args.metallic else None,
                tone=self.args.tone,
                gamma=self.args.gamma,
                occlusion=occlusion.detach(),
                brdf_lut=self.brdf_lut
            )

            diffuse_rgb = pbr_result["diffuse_rgb"].clamp(min=0.0, max=1.0).permute(2, 0, 1)
            diffuse_rgb = torch.where(normal_mask, diffuse_rgb, current_bg[:, None, None])
            render_direct = pbr_result["render_rgb"].permute(2, 0, 1)
            render_direct = torch.where(normal_mask, render_direct, current_bg[:, None, None])

            tanfovx = math.tan(viewpoint_cam.FoVx * 0.5)
            tanfovy = math.tan(viewpoint_cam.FoVy * 0.5)
            SSR = Gaussian_SSR(tanfovx, tanfovy, W, H, self.args.radius, self.args.bias, self.args.thick, self.args.delta, self.args.step, self.args.start)
            
            if self.args.metallic:
                F0 = (1.0 - metallic_map) * 0.04 + albedo_map * metallic_map
            else:
                F0 = torch.ones_like(albedo_map) * 0.04
                metallic_map = torch.zeros_like(roughness_map)
                
            linear_rgb = srgb_to_linear(render_direct)
            (IRR, _) = SSR(out_normal_view.detach(), depth_pos.detach(), linear_rgb.detach(), albedo_map, roughness_map, metallic_map, F0)
            IRR = linear_to_srgb(IRR)
            IRR = kornia.filters.median_blur(IRR[None, ...], (3, 3))[0]
            render_rgb = render_direct + IRR

            loss_target = l1_loss(render_rgb, gt_image)

            if (normal_mask == 0).sum() > 0:
                brdf_tv_loss = get_masked_tv_loss(normal_mask, gt_image, torch.cat([albedo_map, roughness_map, metallic_map], dim=0))
            else:
                brdf_tv_loss = get_tv_loss(gt_image, torch.cat([albedo_map, roughness_map, metallic_map], dim=0), pad=1, step=1)

            lamb_loss = (1.0 - roughness_map[normal_mask]).mean() + metallic_map[normal_mask].mean()
            envmap = dr.texture(self.cubemap.base[None, ...], self.envmap_dirs[None, ...].contiguous(), filter_mode="linear", boundary_mode="cube")[0]
            tv_h1 = torch.pow(envmap[1:, :, :] - envmap[:-1, :, :], 2).mean()
            tv_w1 = torch.pow(envmap[:, 1:, :] - envmap[:, :-1, :], 2).mean()
            env_tv_loss = tv_h1 + tv_w1

            loss_parasitic = brdf_tv_loss * self.args.brdf_tv + lamb_loss * 0.001 + env_tv_loss * self.args.env_tv
            loss = loss_target + loss_parasitic

        grad_cos_sim = 0.0
        parasitic_ratio = 0.0
        stats = {}

        if self.track_decoupling and step % 100 == 0:
            self.gaussians.optimizer.zero_grad(set_to_none=True)
            loss_target.backward(retain_graph=True)
            grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz)
            
            self.gaussians.optimizer.zero_grad(set_to_none=True)
            loss_parasitic.backward(retain_graph=True)
            grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz)
            
            valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0)
            if valid_mask.any():
                grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean())
            parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7))
            
            param_groups_map = {
                "spatial": [self.gaussians._xyz],
                "geometry": [self.gaussians._scaling, self.gaussians._rotation],
                "opacity": [self.gaussians._opacity],
                "appearance": [self.gaussians._features_dc, self.gaussians._features_rest, self.cubemap.base],
            }

            self.gaussians.optimizer.zero_grad(set_to_none=True)
            if step > self.args.pbr_iteration:
                self.light_optimizer.zero_grad(set_to_none=True)

            loss_target.backward(retain_graph=True)
            grads_target = {}
            for group_name, params in param_groups_map.items():
                grads_target[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p is not None and p.grad is not None])

            self.gaussians.optimizer.zero_grad(set_to_none=True)
            if step > self.args.pbr_iteration:
                self.light_optimizer.zero_grad(set_to_none=True)

            loss_parasitic.backward(retain_graph=True)
            grads_parasitic = {}
            for group_name, params in param_groups_map.items():
                grads_parasitic[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p is not None and p.grad is not None])

            for group_name in param_groups_map:
                gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name)
                if gt is not None and gp is not None and gt.numel() > 0 and gp.numel() > 0 and gt.norm() > 0 and gp.norm() > 0:
                    cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0)))
                    r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7))
                    ti = r * max(0.0, -cos)
                else:
                    ti = 0.0
                stats[f"sti_{group_name}"] = ti

            self.gaussians.optimizer.zero_grad(set_to_none=True)
            if step > self.args.pbr_iteration:
                self.light_optimizer.zero_grad(set_to_none=True)
            loss.backward()
        else:
            loss.backward()

        with torch.no_grad():
            if step < self.opt.densify_until_iter:
                self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
                self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
                
                if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0:
                    size_threshold = 20 if step > self.opt.opacity_reset_interval else None
                    self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.05, self.scene.cameras_extent, size_threshold)
                
                if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter):
                    self.gaussians.reset_opacity()

            self.gaussians.optimizer.step()
            self.gaussians.optimizer.zero_grad(set_to_none=True)

            if step >= self.args.pbr_iteration:
                self.light_optimizer.step()
                self.light_optimizer.zero_grad(set_to_none=True)
                self.cubemap.clamp_(min=0.0)

        num_gaussians = self.gaussians.get_xyz.shape[0]
        metrics = {
            "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic),
            "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians),
            "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)),
            "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio),
            "cubemap_energy": float(self.cubemap.base.norm(p=2)),
            "roughness_mean": float(roughness_map.mean()),
            "metallic_ratio": float((metallic_map > 0.5).float().mean())
        }
        metrics.update(stats)
        self.last_n_gaussians = num_gaussians
        
        histograms = {}
        if step % 1000 == 0:
            histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach()
            scales = torch.exp(self.gaussians._scaling).clone().detach()
            histograms["scaling"] = scales
            scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales
            gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7)
            histograms["anisotropy"] = gamma
            histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1)

        return metrics, histograms
        
    def render(self, camera):
        with torch.no_grad():
            render_pkg = render(camera, self.gaussians, self.pipe, self.background, pad_normal=False, derive_normal=True, radius=self.args.radius, bias=self.args.bias, thick=self.args.thick, delta=self.args.delta, step=self.args.step, start=self.args.start)
        return {"image": render_pkg["render"], "depth": render_pkg.get("depth_map", None)}

    def save(self, save_dir, step):
        self.scene.save(step)
        torch.save({"cubemap": self.cubemap.state_dict(), "light_optimizer": self.light_optimizer.state_dict(), "iteration": step}, os.path.join(save_dir, f"chkpnt{step}.pth"))

    def load(self, model_path, iteration):
        self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply'))
        chkpnt = os.path.join(model_path, f"chkpnt{iteration}.pth")
        if os.path.exists(chkpnt):
            ckpt = torch.load(chkpnt)
            self.cubemap.load_state_dict(ckpt["cubemap"])

    def get_spatial_centers(self):
        return self.gaussians._xyz

    def compute_physical_metrics(self, cameras=None):
        metrics = {}
        with torch.no_grad():
            raw_scales = self.gaussians._scaling
            scales = torch.exp(raw_scales)
            scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2)
            max_S, _ = torch.max(scales_2d, dim=1)
            min_S, _ = torch.min(scales_2d, dim=1)
            gamma = max_S / (min_S + 1e-7)
            metrics["gamma_median"] = float(torch.median(gamma))
            metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90))
            metrics["scale_mean"] = float(torch.mean(scales_2d))
            metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity)))
            dc, rest = self.gaussians._features_dc, self.gaussians._features_rest
            if rest is not None and rest.shape[1] > 0:
                metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7))
            if cameras is not None and len(cameras) > 0:
                view_dirs = []
                for c in cameras:
                    view_dirs.append(c.world_view_transform[:3, 2].tolist())
                view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1)
                rots = F.normalize(self.gaussians._rotation.clone(), dim=1)
                w, x, y, z = rots.unbind(dim=-1)
                normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1)
                max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1)
                metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean())
        return metrics

    def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor:
        with torch.no_grad():
            V = query_points.shape[0]
            densities = torch.zeros(V, device="cuda")
            xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze()
            scales = torch.exp(self.gaussians._scaling)
            sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2)
            N_gaussians = xyz.shape[0]
            chunk_size = max(1, 30_000_000 // (N_gaussians + 1))
            for i in range(0, V, chunk_size):
                end = min(i + chunk_size, V)
                dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2)
                weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7))
                densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1)
            return densities