import os
import sys
import random
import torch
import numpy as np
import torch.nn.functional as F
from argparse import ArgumentParser

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

def force_unload_conflicts():
    conflicts = ['render', 'scene', 'utils', 'arguments', 'model']
    for name in list(sys.modules.keys()):
        parts = name.split('.')
        if parts[0] in conflicts:
            del sys.modules[name]

@register_method("bgtriangle")
class BGTriangleWrapper(BaseMethod):
    def __init__(self, dataset_config, hyperparams):
        force_unload_conflicts()
        
        self.method_path = os.path.abspath("/root/autodl-tmp/bg_triangle_official")
        if self.method_path not in sys.path:
            sys.path.insert(0, self.method_path)
        
        try:
            from utils.loss_utils import l1_loss, ssim
            from arguments import ModelParams, PipelineParams, OptimizationParams
            from model.bprimitive_bezier import BPrimitiveBezier
            from render.renderer import Renderer
            from scene import Scene
            
            self.l1_loss = l1_loss
            self.ssim = ssim
            
            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.eval = True
            self.args.resolution = dataset_config.get("resolution", 1)
            self.track_decoupling = hyperparams.get("track_decoupling", False)

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

            self.order = hyperparams.get("order", 1)
            self.bprimitive = BPrimitiveBezier(self.order, self.dataset.sh_degree)
            self.bprimitive.boundary_mode = 1
            
            # 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.bprimitive)
            self.bprimitive.training_setup(self.opt)

            N = self.bprimitive.control_point.shape[0]
            if not hasattr(self.bprimitive, 'vis_map'):
                self.bprimitive.vis_map = torch.zeros((N, 6, 6), device="cuda")
            if not hasattr(self.bprimitive, 'vis_accum'):
                self.bprimitive.vis_accum = torch.zeros((N, 1), device="cuda")
            if not hasattr(self.bprimitive, 'denom_vis'):
                self.bprimitive.denom_vis = torch.zeros((N, 1), device="cuda")
            if not hasattr(self.bprimitive, 'edge_accum'):
                self.bprimitive.edge_accum = torch.zeros((N, 1), device="cuda")
            if not hasattr(self.bprimitive, 'denom_edge'):
                self.bprimitive.denom_edge = torch.zeros((N, 1), device="cuda")
            if not hasattr(self.bprimitive, 'gradient_accum'):
                self.bprimitive.gradient_accum = torch.zeros((N, 1), device="cuda")
            if not hasattr(self.bprimitive, 'denom'):
                self.bprimitive.denom = torch.zeros((N, 1), device="cuda")

            self.renderer = Renderer()
            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_primitives = self.bprimitive.control_point.shape[0]
            
        finally:
            if self.method_path in sys.path:
                sys.path.remove(self.method_path)

    def train_iteration(self, step):
        force_unload_conflicts()
        if self.method_path not in sys.path: 
            sys.path.insert(0, self.method_path)
            
        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))

        render_pkg, debug_info = self.renderer(
            self.bprimitive, viewpoint_cam, self.background, None,
            self.pipe.num_segments_per_bprimitive_edge, -5.0, 1.0, "version_1"
        )
        image = render_pkg
        
        gt_image = viewpoint_cam.original_image.cuda()
        loss_l1 = self.l1_loss(image, gt_image)
        loss_ssim = 1.0 - self.ssim(image, gt_image)
        loss_target = (1.0 - 0.2) * loss_l1
        loss_parasitic = 0.2 * loss_ssim
        loss = loss_target + loss_parasitic

        grad_cos_sim = 0.0
        parasitic_ratio = 0.0

        if self.track_decoupling and step % 100 == 0:
            self.bprimitive.optimizer.zero_grad(set_to_none=True)
            loss_target.backward(retain_graph=True)
            # 使用 reshape 替换 view
            grad_target = self.bprimitive.control_point.grad.clone().reshape(self.bprimitive.control_point.shape[0], -1) if self.bprimitive.control_point.grad is not None else torch.zeros_like(self.bprimitive.control_point).reshape(self.bprimitive.control_point.shape[0], -1)
            
            self.bprimitive.optimizer.zero_grad(set_to_none=True)
            loss_parasitic.backward(retain_graph=True)
            # 使用 reshape 替换 view
            grad_parasitic = self.bprimitive.control_point.grad.clone().reshape(self.bprimitive.control_point.shape[0], -1) if self.bprimitive.control_point.grad is not None else torch.zeros_like(self.bprimitive.control_point).reshape(self.bprimitive.control_point.shape[0], -1)
            
            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))
            
            self.bprimitive.optimizer.zero_grad(set_to_none=True)
            loss.backward()
        else:
            loss.backward()

        with torch.no_grad():
            self.bprimitive.optimizer.step()
            self.bprimitive.optimizer.zero_grad(set_to_none=True)

        num_primitives = self.bprimitive.control_point.shape[0]
        metrics = {
            "loss": float(loss),
            "loss_l1": float(loss_target),
            "loss_ssim": float(loss_parasitic),
            "num_gaussians": int(num_primitives),
            "delta_N": int(num_primitives - self.last_n_primitives),
            "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)),
            "grad_cos_sim": float(grad_cos_sim),
            "parasitic_ratio": float(parasitic_ratio)
        }
        self.last_n_primitives = num_primitives

        histograms = {}
        if step % 1000 == 0:
            histograms["opacity"] = torch.sigmoid(self.bprimitive.opacity).clone().detach()
            scales = torch.exp(self.bprimitive.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.bprimitive.control_point_dc.detach().norm(dim=-1).mean(dim=1)
        
        if self.method_path in sys.path: 
            sys.path.remove(self.method_path)
            
        return metrics, histograms

    def render(self, camera):
        force_unload_conflicts()
        if self.method_path not in sys.path: 
            sys.path.insert(0, self.method_path)
            
        with torch.no_grad():
            image, _ = self.renderer(
                self.bprimitive, camera, self.background, None, 
                self.pipe.num_segments_per_bprimitive_edge, -5.0, 1.0, "version_1"
            )
            
        if self.method_path in sys.path: 
            sys.path.remove(self.method_path)
            
        return {"image": image, "depth": None}

    def save(self, save_dir, step):
        force_unload_conflicts()
        if self.method_path not in sys.path: 
            sys.path.insert(0, self.method_path)
            
        self.scene.save(step)
        
        if self.method_path in sys.path: 
            sys.path.remove(self.method_path)

    def load(self, model_path, iteration):
        pass

    def get_spatial_centers(self):
        return self.bprimitive.control_point.mean(dim=1)

    def compute_physical_metrics(self, cameras=None):
        metrics = {}
        with torch.no_grad():
            raw_scales = self.bprimitive.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.bprimitive.opacity)))
            
            dc = self.bprimitive.control_point_dc
            rest = self.bprimitive.control_point_rest
            if rest is not None and rest.shape[2] > 0:
                metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7))

        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 = self.bprimitive.control_point.mean(dim=1)
            opacities = torch.sigmoid(self.bprimitive.opacity).squeeze()
            scales = torch.exp(self.bprimitive.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