from functools import partial
from typing import Callable

import torch

from gaussian_splatting.trainer.densifier import SplitCloneDensifier, AbstractDensifier
from feature_3dgs import SemanticGaussianModel, FeatureCameraDataset

from .trainer import SemanticDensificationInstruct, SemanticDensificationTrainer


class SemanticSplitCloneDensifier(SplitCloneDensifier):
    """SplitCloneDensifier that also handles _encoded_semantics."""

    @property
    def model(self) -> SemanticGaussianModel:
        return super().model

    def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):
        base = super().densify_and_split(grads, grad_threshold, scene_extent, N)
        # base.replace_xyz_mask is exactly the selected_pts_mask
        selected_pts_mask = base.replace_xyz_mask
        new_encoded_semantics = self.model._encoded_semantics[selected_pts_mask]
        return SemanticDensificationInstruct(
            **base._asdict(),
            new_encoded_semantics=new_encoded_semantics,
        )

    def densify_and_clone(self, grads, grad_threshold, scene_extent):
        selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)
        selected_pts_mask = torch.logical_and(
            selected_pts_mask,
            torch.max(self.model.get_scaling, dim=1).values <= self.densify_percent_dense * scene_extent,
        )

        return SemanticDensificationInstruct(
            new_xyz=self.model._xyz[selected_pts_mask],
            new_features_dc=self.model._features_dc[selected_pts_mask],
            new_features_rest=self.model._features_rest[selected_pts_mask],
            new_opacity=self.model._opacity[selected_pts_mask],
            new_scaling=self.model._scaling[selected_pts_mask],
            new_rotation=self.model._rotation[selected_pts_mask],
            new_encoded_semantics=self.model._encoded_semantics[selected_pts_mask],
        )


# ======================================================================
# Convenience wrapper functions (mirror the base module's pattern)
# ======================================================================

def SemanticSplitCloneDensifierWrapper(
        base_densifier_constructor: Callable[..., AbstractDensifier],
        model: SemanticGaussianModel,
        dataset: FeatureCameraDataset,
        *args,
        densify_from_iter=500,
        densify_until_iter=15000,
        densify_interval=100,
        densify_grad_threshold=0.0002,
        densify_percent_dense=0.01,
        densify_percent_too_big=0.8,
        densify_limit_n=None,
        **configs):
    return SemanticSplitCloneDensifier(
        base_densifier_constructor(model, dataset, *args, **configs),
        dataset,
        densify_from_iter=densify_from_iter,
        densify_until_iter=densify_until_iter,
        densify_interval=densify_interval,
        densify_grad_threshold=densify_grad_threshold,
        densify_percent_dense=densify_percent_dense,
        densify_percent_too_big=densify_percent_too_big,
        densify_limit_n=densify_limit_n
    )


def SemanticSplitCloneDensifierTrainerWrapper(
        base_densifier_constructor: Callable[..., AbstractDensifier],
        model: SemanticGaussianModel, dataset: FeatureCameraDataset, *args,
        **configs):
    return SemanticDensificationTrainer.from_densifier_constructor(
        partial(SemanticSplitCloneDensifierWrapper, base_densifier_constructor),
        model, dataset, *args,
        **configs,
    )
