import numpy as np
import torch
import torch.nn as nn
import json
from diffusers import DDIMScheduler

from diffusionGS.utils.typing import *
from tqdm import tqdm
import math
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
from diffusionGS.utils.structure import Mesh
import inspect
from PIL import Image, ImageOps
import cv2
def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac):
    betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
    warmup_time = int(num_diffusion_timesteps * warmup_frac)
    betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64)
    return betas

def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
    """
    This is the deprecated API for creating beta schedules.
    See get_named_beta_schedule() for the new library of schedules.
    """
    if beta_schedule == "quad":
        betas = (
            np.linspace(
                beta_start ** 0.5,
                beta_end ** 0.5,
                num_diffusion_timesteps,
                dtype=np.float64,
            )
            ** 2
        )
    elif beta_schedule == "linear":
        betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
    elif beta_schedule == "scaled_linear":
        betas = np.linspace(beta_start**0.5, beta_end**0.5, num_diffusion_timesteps, dtype=np.float64)**2
    elif beta_schedule == "warmup10":
        betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1)
    elif beta_schedule == "warmup50":
        betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5)
    elif beta_schedule == "const":
        betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
    elif beta_schedule == "jsd":  # 1/T, 1/(T-1), 1/(T-2), ..., 1
        betas = 1.0 / np.linspace(
            num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64
        )
    else:
        raise NotImplementedError(beta_schedule)
    assert betas.shape == (num_diffusion_timesteps,)
    return betas

def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
    """
    Create a beta schedule that discretizes the given alpha_t_bar function,
    which defines the cumulative product of (1-beta) over time from t = [0,1].
    :param num_diffusion_timesteps: the number of betas to produce.
    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
                      produces the cumulative product of (1-beta) up to that
                      part of the diffusion process.
    :param max_beta: the maximum beta to use; use values lower than 1 to
                     prevent singularities.
    """
    betas = []
    for i in range(num_diffusion_timesteps):
        t1 = i / num_diffusion_timesteps
        t2 = (i + 1) / num_diffusion_timesteps
        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
    return np.array(betas)

def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
    """
    Get a pre-defined beta schedule for the given name.
    The beta schedule library consists of beta schedules which remain similar
    in the limit of num_diffusion_timesteps.
    Beta schedules may be added, but should not be removed or changed once
    they are committed to maintain backwards compatibility.
    """
    if schedule_name == "linear":
        # Linear schedule from Ho et al, extended to work for any number of
        # diffusion steps.
        scale = 1000 / num_diffusion_timesteps
        return get_beta_schedule(
            "linear",
            beta_start=scale * 0.0001,
            beta_end=scale * 0.02,
            num_diffusion_timesteps=num_diffusion_timesteps,
        )
    elif schedule_name == "squaredcos_cap_v2":
        return betas_for_alpha_bar(
            num_diffusion_timesteps,
            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
        )
    else:
        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")

def normal_kl(mean1, logvar1, mean2, logvar2):
    """
    Compute the KL divergence between two gaussians.
    Shapes are automatically broadcasted, so batches can be compared to
    scalars, among other use cases.
    """
    tensor = next(
        (
            obj
            for obj in (mean1, logvar1, mean2, logvar2)
            if isinstance(obj, torch.Tensor)
        ),
        None,
    )
    assert tensor is not None, "at least one argument must be a Tensor"

    # Force variances to be Tensors. Broadcasting helps convert scalars to
    # Tensors, but it does not work for torch.exp().
    logvar1, logvar2 = [
        x if isinstance(x, torch.Tensor) else torch.tensor(x, device=tensor.device)
        for x in (logvar1, logvar2)
    ]

    return 0.5 * (
        -1.0
        + logvar2
        - logvar1
        + torch.exp(logvar1 - logvar2)
        + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
    )

def mean_flat(tensor):
    """
    Take the mean over all non-batch dimensions.
    """
    return tensor.mean(dim=list(range(1, len(tensor.shape))))

def approx_standard_normal_cdf(x):
    """
    A fast approximation of the cumulative distribution function of the
    standard normal.
    """
    return 0.5 * (1.0 + torch.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * torch.pow(x, 3))))

def discretized_gaussian_log_likelihood(x, *, means, log_scales):
    """
    Compute the log-likelihood of a Gaussian distribution discretizing to a
    given image.
    :param x: the target images. It is assumed that this was uint8 values,
              rescaled to the range [-1, 1].
    :param means: the Gaussian mean Tensor.
    :param log_scales: the Gaussian log stddev Tensor.
    :return: a tensor like x of log probabilities (in nats).
    """
    assert x.shape == means.shape == log_scales.shape
    centered_x = x - means
    inv_stdv = torch.exp(-log_scales)
    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
    cdf_plus = approx_standard_normal_cdf(plus_in)
    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
    cdf_min = approx_standard_normal_cdf(min_in)
    log_cdf_plus = torch.log(cdf_plus.clamp(min=1e-12))
    log_one_minus_cdf_min = torch.log((1.0 - cdf_min).clamp(min=1e-12))
    cdf_delta = cdf_plus - cdf_min
    log_probs = torch.where(
        x < -0.999,
        log_cdf_plus,
        torch.where(x > 0.999, log_one_minus_cdf_min, torch.log(cdf_delta.clamp(min=1e-12))),
    )
    assert log_probs.shape == x.shape
    return log_probs

def extract_into_tensor(arr, timesteps, broadcast_shape):
    """
    Extract values from a 1-D numpy array for a batch of indices.
    :param arr: the 1-D numpy array.
    :param timesteps: a tensor of indices into the array to extract.
    :param broadcast_shape: a larger shape of K dimensions with the batch
                            dimension equal to the length of timesteps.
    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
    """
    res = torch.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
    while len(res.shape) < len(broadcast_shape):
        res = res[..., None]
    return res + torch.zeros(broadcast_shape, device=timesteps.device)


# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
    scheduler,
    num_inference_steps: Optional[int] = None,
    device: Optional[Union[str, torch.device]] = None,
    timesteps: Optional[List[int]] = None,
    sigmas: Optional[List[float]] = None,
    **kwargs,
):
    r"""
    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.

    Args:
        scheduler (`SchedulerMixin`):
            The scheduler to get timesteps from.
        num_inference_steps (`int`):
            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
            must be `None`.
        device (`str` or `torch.device`, *optional*):
            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
        timesteps (`List[int]`, *optional*):
            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
            `num_inference_steps` and `sigmas` must be `None`.
        sigmas (`List[float]`, *optional*):
            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
            `num_inference_steps` and `timesteps` must be `None`.

    Returns:
        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
        second element is the number of inference steps.
    """
    if timesteps is not None and sigmas is not None:
        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
    if timesteps is not None:
        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
        if not accepts_timesteps:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                f" timestep schedules. Please check whether you are using the correct scheduler."
            )
        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)
    elif sigmas is not None:
        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
        if not accept_sigmas:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                f" sigmas schedules. Please check whether you are using the correct scheduler."
            )
        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)
    else:
        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
        timesteps = scheduler.timesteps
    return timesteps, num_inference_steps


def euler_sample(ddim_scheduler,
                diffusion_model: torch.nn.Module,
                shape: Union[List[int], Tuple[int]],
                cond: torch.FloatTensor,
                steps: int,
                eta: float = 0.0,
                guidance_scale: float = 3.0,
                do_classifier_free_guidance: bool = True,
                generator: Optional[torch.Generator] = None,
                device: torch.device = "cuda:0",
                disable_prog: bool = True):

    assert steps > 0, f"{steps} must > 0."

    # init latents
    bsz = cond.shape[0]
    if do_classifier_free_guidance:
        bsz = bsz // 2

    # latents = torch.randn(
    #     (bsz, *shape),
    #     generator=generator,
    #     device=cond.device,
    #     dtype=cond.dtype,
    # )
    latents = diffusion_model.get_init_point(bsz)
    #latents = latents + init_offsets
    # scale the initial noise by the standard deviation required by the scheduler
    #latents = latents * ddim_scheduler.init_noise_sigma
    # set timesteps
    ddim_scheduler.set_timesteps(steps)
    timesteps = ddim_scheduler.timesteps.to(device)
    # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
    # eta (η) is only used with the DDIMScheduler, and between [0, 1]
    extra_step_kwargs = {
        # "eta": eta,
        "generator": generator
    }
    
    #timesteps, num_inference_steps = retrieve_timesteps(ddim_scheduler, steps, device)
    #breakpoint()
    # reverse
    for i, t in enumerate(tqdm(timesteps, disable=disable_prog, desc="Euler Sampling:", leave=False)):
        # expand the latents if we are doing classifier free guidance
        latent_model_input = (
            torch.cat([latents] * 2)
            if do_classifier_free_guidance
            else latents
        )
        #breakpoint()
        # predict the noise residual
        timestep_tensor = torch.tensor(t.clone().detach(), dtype=torch.long, device=device)
        timestep_tensor = timestep_tensor.expand(latent_model_input.shape[0])
        noise_pred = diffusion_model.forward(latent_model_input, timestep_tensor, cond, train_stage = False)

        # perform guidance
        if do_classifier_free_guidance:
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (
                    noise_pred_text - noise_pred_uncond
            )

        # compute the previous noisy sample x_t -> x_t-1
        latents = ddim_scheduler.step(
            noise_pred, t, latents, **extra_step_kwargs
        ).prev_sample

        yield latents, t


def compute_snr(alphas_cumprod, timesteps):
    """
    Computes SNR as per
    https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
    """
    sqrt_alphas_cumprod = alphas_cumprod**0.5
    sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5

    # Expand the tensors.
    # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026
    sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
    while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
        sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
    alpha = sqrt_alphas_cumprod.expand(timesteps.shape)

    sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
    while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):
        sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]
    sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)

    # Compute SNR.
    snr = (alpha / sigma) ** 2
    return snr

def get_val_data(file_list, n_supervision=4096, geometry_type="occupancies"):
    batch_rand_points = []
    batch_occupancies = []
    for file in file_list:
        points = np.load(file)
        rand_points = np.asarray(points['points']) * 2 # range from -1.1 to 1.1
        rand_points = torch.from_numpy(rand_points)
        rand_points = torch.split(rand_points, n_supervision, dim=0)
        rand_points = torch.stack(rand_points[0:-1])
        occupancies = np.asarray(points[geometry_type])
        occupancies = np.unpackbits(occupancies)
        occupancies = torch.from_numpy(occupancies)
        occupancies = torch.split(occupancies, n_supervision, dim=0)
        occupancies = torch.stack(occupancies[0:-1])
        batch_rand_points.append(rand_points)
        batch_occupancies.append(occupancies)
    batch_rand_points = torch.stack(batch_rand_points)
    batch_occupancies= torch.stack(batch_occupancies)
    B, M, N, _ = batch_rand_points.shape
    return batch_rand_points.view(B*M, N, 3), batch_occupancies.view(B*M, N)

def get_mvp_matrix(c2w:Tensor, proj_mtx:Tensor):
    '''
    c2w: [batch_size, 4, 4]
    proj_mtx: [batch_size, 4, 4]

    mvp_mtx: [batch_size, 4, 4]
    '''
    # calculate w2c from c2w: R' = Rt, t' = -Rt * t
    # mathematically equivalent to (c2w)^-1
    w2c = torch.zeros(c2w.shape[0], 4, 4).to(c2w)
    w2c[:, :3, :3] = c2w[:, :3, :3].permute(0, 2, 1)
    w2c[:, :3, 3:] = -c2w[:, :3, :3].permute(0, 2, 1) @ c2w[:, :3, 3:]
    w2c[:, 3, 3] = 1.0
    # calculate mvp matrix by proj_mtx @ w2c (mv_mtx)
    mvp_mtx = proj_mtx @ w2c
    return mvp_mtx.to(c2w)

def get_projection_matrix_perspective(intrinsics:Tensor, near: float = 0.01, far: float = 100.0):
    '''
    intrinsics: [batch_size, 3, 3], normalized

    proj_mtx: [batch_size, 4, 4]
    '''
    batch_size = intrinsics.shape[0]
    proj_mtx = torch.zeros(batch_size, 4, 4, dtype=torch.float32).to(intrinsics)
    proj_mtx[:, 0, 0] = 2 * intrinsics[:, 0, 0]
    proj_mtx[:, 1, 1] = - 2 * intrinsics[:, 1, 1] # add a negative sign here as the y axis is flipped in nvdiffrast output
    proj_mtx[:, 2, 2] = -(far + near) / (far - near)
    proj_mtx[:, 2, 3] = -2.0 * far * near / (far - near)
    proj_mtx[:, 3, 2] = -1.0
    return proj_mtx


def render_mesh(mesh, cam2world_matrices, intrinsics, device, height=224, width=224, radius=1):
    import nvdiffrast.torch as dr

    ## world to ndc
    v_homo = torch.cat([mesh.v_pos, torch.ones([mesh.v_pos.shape[0], 1]).to(mesh.v_pos)], dim=-1)
    mvp_mtx = get_mvp_matrix(cam2world_matrices, get_projection_matrix_perspective(intrinsics))
    v_pos_clip = torch.matmul(v_homo, mvp_mtx.permute(0, 2, 1)).float()

    v_nrm = mesh.v_nrm
    
    ## rasterize
    ctx = dr.RasterizeCudaContext(device)
    mesh.t_pos_idx = mesh.t_pos_idx.to(torch.int32)
    rast, _ = dr.rasterize(ctx, v_pos_clip, mesh.t_pos_idx, (height, width))
    mask = rast[..., 3:] > 0

    ## render alpha
    mask_aa = dr.antialias(mask.float(), rast, v_pos_clip, mesh.t_pos_idx)

    ## render depth
    gb_depth, _ = dr.interpolate(v_pos_clip[:, :, 3:].contiguous(), rast, mesh.t_pos_idx)
    gb_depth_aa = torch.lerp(torch.zeros_like(gb_depth), gb_depth, mask.float())
    # gb_depth_aa = dr.antialias(gb_depth_aa, rast, v_pos_clip, mesh.t_pos_idx)

    ## render normal
    gb_normal, _ = dr.interpolate(v_nrm, rast, mesh.t_pos_idx)
    gb_normal = F.normalize(gb_normal, dim=-1)
    gb_normal_aa = torch.lerp(torch.full_like(gb_normal, fill_value=-1.0), gb_normal, mask.float())
    gb_normal_aa = dr.antialias(gb_normal_aa, rast, v_pos_clip, mesh.t_pos_idx)

    ## render rgb
    selector = mask[..., 0]
    gb_pos, _ = dr.interpolate(mesh.v_pos, rast, mesh.t_pos_idx)
    positions = gb_pos[selector]
    gb_rgb_fg = torch.zeros(*mask_aa.shape[:3], 3).to(device)
    gb_rgb_fg[selector] = torch.ones(3).to(device) * 125
    gb_rgb_bg = torch.ones_like(gb_rgb_fg).to(gb_rgb_fg)
    gb_rgb = torch.lerp(gb_rgb_bg, gb_rgb_fg, mask.to(gb_rgb_fg))
    gb_rgb_aa = dr.antialias(gb_rgb, rast, v_pos_clip, mesh.t_pos_idx)
    # out["rgb"] = gb_rgb_aa

    out = {
        "alpha": mask_aa,
        "depth": gb_depth_aa,
        "rgb": gb_rgb_aa,
        "normal": gb_normal_aa
    }
    return out

def get_camera(device, img_size=224, focal=5, distance=10):
    cam2world_matrices = torch.as_tensor([
                                [[1, 0, 0, 0],
                                [0, 1, 0, 0],
                                [0, 0, 1, distance],
                                [0, 0, 0, 1]], # front to back

                                [[0, 0, 1, distance],
                                [0, 1, 0, 0],
                                [1, 0, 0, 0],
                                [0, 0, 0, 1]], # front to back

                                [[-1, 0, 0, 0],
                                [0, 1, 0, 0],
                                [0, 0, -1, -distance],
                                [0, 0, 0, 1]], # front to back

                                [[0, 0, -1, -distance],
                                [0, 1, 0, 0],
                                [-1, 0, 0, 0],
                                [0, 0, 0, 1]], # front to back
                            ], dtype=torch.float32).to(device)
    intrinsics = torch.zeros([3, 3]).to(device)
    intrinsics[0, 0] = focal
    intrinsics[1, 1] = focal
    intrinsics[1, 2] = img_size // 2
    intrinsics[0, 2] = img_size // 2
    intrinsics = intrinsics.unsqueeze(0).repeat(4, 1 ,1).to(torch.float32)

    return cam2world_matrices, intrinsics

def read_image(img, img_size=224):
    transform = transforms.Compose(
            [
                transforms.Resize(img_size, transforms.InterpolationMode.BICUBIC, antialias=True),
                transforms.CenterCrop(img_size),  # crop a (224, 224) square
                transforms.ToTensor()
            ]
        )
    rgb = Image.open(img)
    rgb = transform(rgb)[:3,...].permute(1, 2, 0)
    return rgb

def compute_metric(shape_model, sample_inputs, sample_outputs, device):
    threshold = 0
    if len(sample_inputs['occupancy']) == 0:
        return 0, 0
    queries, labels = get_val_data(sample_inputs['occupancy'])
    queries = queries.to(device)
    labels = labels.to(device)
    latent = sample_outputs[0][:len(sample_inputs['occupancy']),...]
    latent = latent.unsqueeze(1).repeat(1, queries.shape[0] // latent.shape[0], 1, 1).view(queries.shape[0], latent.shape[1], latent.shape[2])
    outputs = shape_model.query(queries, latent)
    pred = torch.zeros_like(outputs)
    pred[outputs>=threshold] = 1
    torch.cuda.empty_cache()
    accuracy = (pred==labels).float().sum(dim=1) / labels.shape[1]
    accuracy = accuracy.mean()
    intersection = (pred * labels).sum(dim=1)
    union = (pred + labels).gt(0).sum(dim=1)
    iou = intersection * 1.0 / union + 1e-5
    iou = iou.mean()

    return accuracy, iou

def compute_metric_from_sdf(shape_model, sample_inputs, sample_outputs, threshold, device):
    if len(sample_inputs['sdf']) == 0:
        return 0, 0
    queries, sdf = get_val_data(sample_inputs['sdf'])
    queries = queries.to(device)
    labels = labels.to(device)
    latent = sample_outputs[0][:len(sample_inputs['occupancy']),...]
    latent = latent.unsqueeze(1).repeat(1, queries.shape[0] // latent.shape[0], 1, 1).view(queries.shape[0], latent.shape[1], latent.shape[2])
    outputs = shape_model.query(queries, latent)
    pred = torch.zeros_like(outputs)
    pred[outputs>=threshold] = 1
    torch.cuda.empty_cache()
    accuracy = (pred==labels).float().sum(dim=1) / labels.shape[1]
    accuracy = accuracy.mean()
    intersection = (pred * labels).sum(dim=1)
    union = (pred + labels).gt(0).sum(dim=1)
    iou = intersection * 1.0 / union + 1e-5
    iou = iou.mean()

    return accuracy, iou

def get_render(vertices, faces, device, img_size=224, focal=5, distance=12):
    cam2world_matrices, intrinsics = get_camera(device, img_size, focal, distance)
    mesh = Mesh(
                v_pos=torch.from_numpy(vertices).to(device),
                t_pos_idx=torch.from_numpy(faces).to(device),
            )
    out = render_mesh(mesh, cam2world_matrices, intrinsics, device, img_size, img_size)
    v, h, w, _ = out['normal'].shape
    normal = out['normal'].transpose(1, 0).contiguous().view(h, w*v, 3)
    depth = out['depth'].transpose(1, 0).contiguous().view(h, w*v, 1)
    if depth[depth != 0].shape[0] != 0:
        depth[depth != 0] = (depth[depth != 0] - depth[depth != 0].min()) / (depth[depth != 0].max() - depth[depth != 0].min())

    return normal, depth


def ema_update(model_dest: nn.Module, model_src: nn.Module, rate):
    param_dict_src = dict(model_src.named_parameters())
    for p_name, p_dest in model_dest.named_parameters():
        p_src = param_dict_src[p_name]
        assert p_src is not p_dest
        p_dest.data.mul_(rate).add_((1 - rate) * p_src.data)


def get_bbox(mask:np.ndarray):
    assert mask.ndim == 2
    row = mask.sum(-1)
    col = mask.sum(-2)

    row_idx = np.where(row > 0)[0]
    col_idx = np.where(col > 0)[0]
    x1, y1, x2, y2 = col_idx.min(), row_idx.min(), col_idx.max(), row_idx.max()
    return np.array([x1, y1, x2, y2])


def preprocess(image:Image.Image,session) -> Image.Image:
    H_1 = W_1 = 512
    scale = 0.8
    color = 'white'
    image = ImageOps.exif_transpose(image)
    if image.mode == 'RGBA' and np.sum(np.array(image.getchannel('A')) > 0) < image.size[0] * image.size[1] - 8:
        rgba = image.copy()
    else:
        # https://github.com/pymatting/pymatting/issues/19
        import rembg
        rgba = rembg.remove(image, alpha_matting=True, session=session)

    alpha = rgba.getchannel('A')
    bboxs = get_bbox(np.array(alpha))
    x1, y1, x2, y2 = bboxs
    dy, dx = y2 - y1, x2 - x1
    s = min(H_1 * scale / dy, W_1 * scale / dx)
    Ht, Wt = int(dy * s), int(dx * s)
    ox, oy = int((W_1 - Wt) / 2), int((H_1 - Ht) / 2)
    bboxt = np.array([ox, oy, ox+Wt, oy+Ht])

    rgba = rgba.crop(bboxs).resize((Wt, Ht))
    fg = rgba.convert('RGB')
    alpha = rgba.getchannel('A')
    alphat = Image.new('L', (W_1, H_1))
    alphat.paste(alpha, bboxt)

    inp_1 = Image.new('RGBA', (W_1, H_1), color)
    inp_1.paste(fg, bboxt, alpha)
    inp_1.putalpha(alphat)
    return inp_1


# def preprocess(image:Image.Image,session) -> Image.Image:
#     H_1 = W_1 = 512
#     color = 'white'
#     image = ImageOps.exif_transpose(image)
#     if image.mode == 'RGBA' and np.sum(np.array(image.getchannel('A')) > 0) < image.size[0] * image.size[1] - 8:
#         rgba = image.copy()
#     else:
#         # https://github.com/pymatting/pymatting/issues/19
#         rgba = rembg.remove(image, alpha_matting=True, session=session)


#     fg = rgba.convert('RGB')
#     alpha = rgba.getchannel('A')
#     alphat = Image.new('L', (W_1, H_1))
#     alphat.paste(alpha)

#     inp_1 = Image.new('RGBA', (W_1, H_1), color)
#     inp_1.paste(fg, mask = alpha)
#     inp_1.putalpha(alphat)
#     return inp_1


def TransformInput(image, c2w, fxfycxcy, patch_size=None):
    assert image.dim() == 5, f"image dim should be 5, but got {image.dim()}"        # b, v, c, h, w
    assert (
        fxfycxcy.dim() == 3
    ), f"fxfycxcy dim should be 3, but got {fxfycxcy.dim()}"                # b, v, 4
    assert c2w.dim() == 4, f"c2w dim should be 4, but got {c2w.dim()}"      # b, v, 4, 4

    b, v, c, h, w = image.size()
    image = image.reshape(b * v, c, h * w)
    fxfycxcy = fxfycxcy.reshape(b * v, 4)
    c2w = c2w.reshape(b * v, 4, 4)          # 进来就把 b 和 v 合并

    # y - i - 行 - h, x - j - 列 - w
    # y 和 x 都是形状为 [h, w] 的矩阵, y 的每一行都是 0 到 w-1, x 的每一列都是 0 到 h-1
    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")

    # print(f"y device: {y.device}, x device: {x.device}, image device: {image.device}")
    # print(f"y dtype: {y.dtype}, x dtype: {x.dtype}, image dtype: {image.dtype}")
    # print(f"y shape: {y.shape}, x shape: {x.shape}, image shape: {image.shape}")
    # stx()
    '''
        1. not eval:
            y device: cpu, x device: cpu, image device: cuda:0
            y dtype: torch.int64, x dtype: torch.int64, image dtype: torch.float32
            y shape: torch.Size([256, 256]), x shape: torch.Size([256, 256]), image shape: torch.Size([64, 3, 65536])
        2. eval:
            y device: cpu, x device: cpu, image device: cuda:0
            y dtype: torch.int64, x dtype: torch.int64, image dtype: torch.float32
            y shape: torch.Size([256, 256]), x shape: torch.Size([256, 256]), image shape: torch.Size([64, 3, 65536])
    '''
    y, x = y.to(image.device), x.to(image.device)

    # [h, w], normalize 到 [-1, 1]
    y_norm, x_norm = (y + 0.5) / h * 2 - 1, (x + 0.5) / w * 2 - 1
    # [b, v, 2, h, w]
    xy_norm = torch.stack([x_norm, y_norm], dim=0)[None, None, :, :, :].expand(
        b, v, -1, -1, -1
    )

    x = x[None, :, :].expand(b * v, -1, -1).reshape(b * v, -1)
    y = y[None, :, :].expand(b * v, -1, -1).reshape(b * v, -1)

    # 到此处, x, y 是图像坐标 (u, v), 它与相机坐标 (x_c, y_c, z_c) 的转换关系是:
    # u = f_x * x_c / z_c + cx
    # v = f_y * y_c / z_c + cy
    # 所以反算 x_c, y_c 是
    # x_c / z_c = (u - cx) / f_x, y_c / z_c = (v - cy) / f_y
    # 那么 per-pixel 的射线方向就是 (x_c, y_c, z_c) = (x_c / z_c, y_c / z_c, 1)

    x = (x + 0.5 - fxfycxcy[:, 2:3]) / fxfycxcy[:, 0:1]
    y = (y + 0.5 - fxfycxcy[:, 3:4]) / fxfycxcy[:, 1:2]
    z = torch.ones_like(x)

    # ray_d 是相机坐标系下的射线方向
    ray_d = torch.stack([x, y, z], dim=2)  # [b*v, h*w, 3]
    ray_d_cam = ray_d.clone()
    # bmm = batch matrix multiplication, 这里转成世界坐标系
    ray_d = torch.bmm(ray_d.to(c2w), c2w[:, :3, :3].transpose(1, 2))  # [b*v, h*w, 3]

    # norm
    ray_d = ray_d / torch.norm(ray_d, dim=2, keepdim=True)  # [b*v, h*w, 3]
    ray_d_cam = ray_d_cam / torch.norm(ray_d_cam, dim=2, keepdim=True)
    
    ray_o = c2w[:, :3, 3][:, None, :].expand_as(ray_d)  # [b*v, h*w, 3]

    # fetch color, ray_o, ray_d for all patch centers
    ray_color_patch, ray_o_patch, ray_d_patch = None, None, None
    ray_xy_norm_patch, proj_mat = None, None
    if patch_size is not None:
        start_patch_center = patch_size / 2.0
        y, x = torch.meshgrid(
            torch.arange(h // patch_size) * patch_size + start_patch_center,
            torch.arange(w // patch_size) * patch_size + start_patch_center,
            indexing="ij",
        )
        y, x = y.to(image.device), x.to(image.device)
        x = x[None, :, :].expand(b * v, -1, -1).reshape(b * v, -1)
        y = y[None, :, :].expand(b * v, -1, -1).reshape(b * v, -1)

        ray_xy_norm_patch = torch.stack(
            [x / w, y / h], dim=2
        )  # use [0,1] for patch center
        K_norm = (
            torch.eye(3, device=image.device).unsqueeze(0).repeat(b * v, 1, 1)
        )  # [b*v, 3, 3]
        K_norm[:, 0, 0] = fxfycxcy[:, 0] / w
        K_norm[:, 1, 1] = fxfycxcy[:, 1] / h
        K_norm[:, 0, 2] = fxfycxcy[:, 2] / w
        K_norm[:, 1, 2] = fxfycxcy[:, 3] / h
        w2c = torch.inverse(c2w)  # [b*v, 4, 4]
        proj_mat = torch.bmm(K_norm.to(w2c), w2c[:, :3, :4])  # [b*v, 3, 4]
        proj_mat = proj_mat.reshape(b * v, 12)
        proj_mat = proj_mat / (proj_mat.norm(dim=1, keepdim=True) + 1e-6)
        proj_mat = proj_mat.reshape(b * v, 3, 4)
        proj_mat = proj_mat * proj_mat[:, 0:1, 0:1].sign()

        # fetch colors for all patch centers
        ray_color_patch = (
            F.grid_sample(
                image.reshape(b * v, c, h, w),
                torch.stack([x / w * 2.0 - 1.0, y / h * 2.0 - 1.0], dim=2).reshape(
                    b * v, -1, 1, 2
                ),
                align_corners=False,
            )
            .squeeze(-1)
            .permute(0, 2, 1)
        ).contiguous()  # [b*v, h'*w', c]

        x = (x - fxfycxcy[:, 2:3]) / fxfycxcy[:, 0:1]
        y = (y - fxfycxcy[:, 3:4]) / fxfycxcy[:, 1:2]
        z = torch.ones_like(x)
        ray_d_patch = torch.stack([x, y, z], dim=2)  # [b*v, h'*w', 3]
        ray_d_patch = torch.bmm(
            ray_d_patch.to(c2w), c2w[:, :3, :3].transpose(1, 2)
        )  # [b*v, h'*w', 3]
        ray_d_patch = ray_d_patch / torch.norm(
            ray_d_patch, dim=2, keepdim=True
        )  # [b*v, h'*w', 3]
        ray_o_patch = c2w[:, :3, 3][:, None, :].expand_as(
            ray_d_patch
        )  # [b*v, h'*w', 3]

        n_patch = ray_color_patch.size(1)
        ray_color_patch = ray_color_patch.reshape(b, v, n_patch, c)
        ray_o_patch = ray_o_patch.reshape(b, v, n_patch, 3)
        ray_d_patch = ray_d_patch.reshape(b, v, n_patch, 3)
        ray_xy_norm_patch = ray_xy_norm_patch.reshape(b, v, n_patch, 2)
        proj_mat = proj_mat.reshape(b, v, 3, 4)

    ray_o = ray_o.reshape(b, v, h, w, 3).permute(0, 1, 4, 2, 3)
    ray_d = ray_d.reshape(b, v, h, w, 3).permute(0, 1, 4, 2, 3)
    ray_d_cam = ray_d_cam.reshape(b, v, h, w, 3).permute(0, 1, 4, 2, 3)
    image = image.reshape(b, v, c, h, w)
    fxfycxcy = fxfycxcy.reshape(b, v, 4)
    c2w = c2w.reshape(b, v, 4, 4)
    return ray_o, ray_d


# 根据 timestep 把 x_t 做成 video 的可视化代码
def display_timestep_on_video(frames, timesteps):
    """
    frames: [T, H, W, C]
    timesteps: [T, ]
    """
    assert len(frames) == len(
        timesteps
    ), f"len(frames) {len(frames)} != len(timesteps) {len(timesteps)}"

    dtype = frames.dtype
    if dtype == np.float32:
        frames = (frames * 255.0).clip(0.0, 255.0).astype(np.uint8)

    # Use putText() method for
    # inserting text on video
    for i in range(len(frames)):
        # stx()
        frames[i] = cv2.putText(
            frames[i],
            f"t={timesteps[i]}",
            (50, 50),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 255, 255),
            2,
            cv2.LINE_4,
        )

    if dtype == np.float32:
        frames = frames.astype(np.float32) / 255.0

    return frames