import os
import glob
import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from body_model import OP_IGNORE_JOINTS
from geometry.mesh import save_mesh_scenes, vertices_to_trimesh
from util.logger import Logger, log_cur_stats
from util.tensor import move_to, detach_all
from vis.output import prep_result_vis, animate_scene

from .losses import RootLoss, SMPLLoss #, MotionLoss
from .output import save_camera_json
import time

LINE_SEARCH = "strong_wolfe"
LIGHT_BLUE=(0.65098039,  0.74117647,  0.85882353)

"""
Optimization happens in multiple stages
Stage 1: fit root orients and trans of each frame indendpently 
Stage 2: fit SMPL poses and betas of each frame independently
Stage 3: fit poses and roots in same global coordinate frame
"""


class StageOptimizer(object):
    def __init__(
        self,
        name,
        model,
        param_names,
        lr=1.0,
        lbfgs_max_iter=20,
        save_every=10,
        vis_every=-1,
        save_meshes=True,
        max_chunk_steps=10,
        **kwargs,
    ):
        Logger.log(f"INITIALIZING OPTIMIZER {name} for {param_names}")
        self.name = name
        self.model = model

        self.set_opt_vars(param_names)

        self.optim = torch.optim.LBFGS(
            self.opt_params, max_iter=lbfgs_max_iter, lr=lr, line_search_fn=LINE_SEARCH
        )
        # LBFGS computes losses multiple times per iteration,
        # save a dict mapping iteration to list of stats_dicts
        self.loss_dicts = {}

        self.save_every = save_every
        self.vis_every = vis_every
        self.save_meshes = save_meshes
        self.max_chunk_steps = max_chunk_steps

        self.cur_step = 0

        self.add_chunk = 0
        self.cur_loss = 0
        self.prev_loss = np.inf
        self.last_updated = 0
        self.reached_max = False
        self.reached_max_iter = -1

    def set_opt_vars(self, param_names):
        Logger.log("Set param names:")
        Logger.log(param_names)

        self.param_names = param_names
        self.model.params.set_require_grads(self.param_names)
        self.opt_params = [
            getattr(self.model.params, name) for name in self.param_names
        ]

    def forward_pass(self, obs_data):
        raise NotImplementedError

    def load_checkpoint(self, out_dir, device=None):
        if device is None:
            device = torch.device("cpu")

        param_path = os.path.join(out_dir, f"{self.name}_params.pth")
        if os.path.isfile(param_path):
            param_dict = torch.load(param_path, map_location=device)
            self.model.params.load_dict(param_dict)
            Logger.log(f"Params loaded from {param_path}")

        optim_path = os.path.join(out_dir, f"{self.name}_optim.pth")
        if os.path.isfile(optim_path):
            optim_dict = torch.load(optim_path)
            self.optim.load_state_dict(optim_dict["optim"])
            self.cur_step = optim_dict["cur_step"]
            Logger.log(f"Optimizer loaded from {optim_path} at iter {self.cur_step}")

    def save_checkpoint(self, out_dir):
        param_path = os.path.join(out_dir, f"{self.name}_params.pth")
        param_dict = self.model.params.get_dict()
        if "world_scale" in param_dict:
            print("WORLD_SCALE", param_dict["world_scale"].detach().cpu())
        if "floor_plane" in param_dict:
            print("FLOOR PLANE", param_dict["floor_plane"].detach().cpu())
        if "cam_f" in param_dict:
            print("CAM_F", param_dict["cam_f"].detach().cpu())
        torch.save(param_dict, param_path)
        Logger.log(f"Model saved at {param_path}")

        optim_path = os.path.join(out_dir, f"{self.name}_optim.pth")
        torch.save(
            {"optim": self.optim.state_dict(), "cur_step": self.cur_step},
            optim_path,
        )
        Logger.log(f"Optimizer saved at {optim_path}")

    def save_results(self, out_dir, seq_name):
        """
        pred dict will be a dictionary of trajectories.
        each trajectory will have params and lists of trimesh sequences
        """
        os.makedirs(out_dir, exist_ok=True)

        with torch.no_grad():
            print('go into get_optim_result from optimizer.py save_results')
            pred_dict = self.model.get_optim_result()
        pred_dict = move_to(detach_all(pred_dict), "cpu")

        i = self.cur_step
        for name, results in pred_dict.items():
            print('save_results, name: ', name)
            # save parameters of trajectory
            out_path = f"{out_dir}/{seq_name}_{i:06d}_{name}_results.npz"
            Logger.log(f"saving params to {out_path}")
            np.savez(out_path, **results)
            print('cam_R in save_results and its .npz', results['cam_R'])

        # also save the cameras
        # print('save the cameras in ptimizer.py save_results')
        # print(self.model.params)
        # with torch.no_grad():
        #     cam_R, cam_t = self.model.params.get_extrinsics()
        #     intrins = self.model.params.get_intrinsics()
        # print('cam_R in 0000.json:', cam_R)
        # save_camera_json(
        #     f"{out_dir}/{seq_name}_cameras_{self.cur_step:06d}.json",
        #     cam_R.detach().cpu(),
        #     cam_t.detach().cpu(),
        #     intrins.detach().cpu(),
        # )

        # plot losses
        self.plot_losses(out_dir)

    # def save_meshes_all(self, res_dir, obs_data, seq_name, num_steps=-1):

    #     # check which results are saved
    #     seq_name = obs_data["seq_name"][0]
    #     res_pre = f"{res_dir}/{seq_name}_opt_{self.cur_step:06d}"
    #     with torch.no_grad():
    #         print('go into get_optim_result from optimizer.py vis_result')
    #         pred_dict = self.model.get_optim_result(num_steps=num_steps)

    #     res_dict = detach_all(pred_dict["world"])
    #     scene_dict = move_to(
    #         prep_result_vis(
    #             res_dict,
    #             obs_data["vis_mask"],
    #             obs_data["track_id"],
    #             self.model.body_model,
    #             temporal_smooth=True
    #         ),
    #         "cpu",
    #     )

    #     # save the meshes
    #     os.makedirs(res_dir, exist_ok=True)
    #     i = self.cur_step
    #     scene_dir = f"{res_dir}/{seq_name}_{i:06d}_meshes"
    #     os.makedirs(scene_dir, exist_ok=True)

    #     verts, joints, colors, l_faces, r_faces, is_right, bounds = scene_dict["geometry"]
    #     T = len(verts)
    #     print(f"{T} mesh frames")
    #     times = list(range(0, T, 1))
    #     for t in times:
    #         if len(is_right[t]) > 1:
    #             assert (is_right[t].cpu().numpy().tolist() == [0,1])
    #             # l_meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
    #             # r_meshes = make_batch_mesh(verts[t][1][None], r_faces[t], colors[t][1][None])
    #             # assert len(l_meshes) == 1
    #             # assert len(r_meshes) == 1
    #             tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
    #             tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))

    #             tmesh = vertices_to_trimesh(verts[t][1].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
    #             tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))

    #         else:
    #             assert len(is_right[t]) == 1
    #             if is_right[t] == 0:
    #                 # meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
    #                 tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
    #                 tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
    #             elif is_right[t] == 1:
    #                 # meshes = make_batch_mesh(verts[t][0][None], r_faces[t], colors[t][0][None])
    #                 tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
    #                 tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))

    def vis_result(self, res_dir, obs_data, vis=None, num_steps=-1):
        if vis is None or self.vis_every < 0:
            return

        # check which results are saved
        seq_name = obs_data["seq_name"][0]
        res_pre = f"{res_dir}/{seq_name}_opt_{self.cur_step:06d}"
        with torch.no_grad():
            print('go into get_optim_result from optimizer.py vis_result')
            pred_dict = self.model.get_optim_result(num_steps=num_steps)

        res_dict = detach_all(pred_dict["world"])
        scene_dict = move_to(
            prep_result_vis(
                res_dict,
                obs_data["vis_mask"],
                obs_data["track_id"],
                self.model.body_model,
            ),
            "cpu",
        )
        animate_scene(vis, scene_dict, res_pre, render_views=["src_cam", "above"])

    def log_losses(self, stats_dict):
        stats_dict = move_to(detach_all(stats_dict), "cpu")
        log_cur_stats(
            stats_dict,
            iter=self.cur_step,
            to_stdout=(self.cur_step % self.save_every == 0),
        )
        for loss_name, loss_val in stats_dict.items():
            loss_dict = self.loss_dicts.get(loss_name, {})
            loss_series = loss_dict.get(self.cur_step, [])
            loss_series.append(loss_val)
            loss_dict[self.cur_step] = loss_series
            self.loss_dicts[loss_name] = loss_dict

    def record_current_losses(self, writer):
        """
        record the mean of current step's loss values in tensorboard
        """
        if len(self.loss_dicts) < 1:
            return

        for loss_name, loss_dict in self.loss_dicts.items():
            loss_mean = np.mean(loss_dict[self.cur_step])
            writer.add_scalar(f"{self.name}/{loss_name}", loss_mean, self.cur_step)

    def plot_losses(self, res_dir):
        """
        plot a box plot for each BFGS iteration
        """
        if len(self.loss_dicts) < 1:
            return
        for loss_name, loss_dict in self.loss_dicts.items():
            # times (list len T)
            # loss vals (list len T of loss value lists)
            times, loss_vals = zip(*loss_dict.items())
            plt.figure()
            plt.boxplot(loss_vals, labels=times, showfliers=False)
            plt.savefig(f"{res_dir}/{loss_name}.png")

    def run(self, obs_data, num_iters, out_dir, vis=None, writer=None):
        self.cur_step = 0
        self.loss.cur_step = 0
        res_dir = os.path.join(out_dir, self.name)
        os.makedirs(res_dir, exist_ok=True)
        seq_name = obs_data["seq_name"][0]
        print("SEQ NAME", seq_name, "num_iters", num_iters)

        # try to load from checkpoint if exists
        device = obs_data["joints2d"].device
        self.load_checkpoint(out_dir, device=device)

        if self.cur_step >= num_iters:
            Logger.log(f"Current optimizer {self}")
            Logger.log(f"Checkpoint at {self.cur_step} >= {num_iters}, skipping")
            return

        Logger.log(f"OPTIMIZING {self.name} FOR {num_iters} ITERATIONS")
        # time.sleep(5)

        # save initial results and vis
        print('go into save_results in optimizers.py run()')
        self.save_results(res_dir, seq_name)

        for i in range(self.cur_step, num_iters):
            Logger.log("ITER: %d" % (i))

            if (i + 1) % self.save_every == 0:  # save before
                self.save_checkpoint(out_dir)
                self.save_results(res_dir, seq_name)
            else:
                self.save_checkpoint(out_dir)

            if (i + 1) % self.vis_every == 0:  # render
                self.vis_result(res_dir, obs_data, vis)

            self.cur_step = i
            self.loss.cur_step = i

            #print("ITER: %d" % (i))
            self.optim_step(obs_data, i, writer)

            # early termination in case of nans
            if np.isnan(self.cur_loss):
                # we need to backtrack
                self.load_checkpoint(out_dir, device=device)
                print(self.cur_loss, self.kkk, self.ppp)
                raise ValueError

            # termination for the last chunk
            if self.reached_max and self.reached_max_iter < 0:
                self.reached_max_iter = i - 1
            if self.reached_max and i - self.reached_max_iter >= self.max_chunk_steps:
                break

            # termination for middle chunks
            loss_change = self.prev_loss - self.cur_loss
            if self.last_updated == i - 1 and loss_change == 0:
                break
            if (
                (self.cur_loss < 0 and loss_change < 100)
                or (i - self.last_updated >= self.max_chunk_steps)
                or (loss_change < 20 and i - self.last_updated > 5)
            ):
                self.add_chunk = self.add_chunk + 1
                self.last_updated = i
            self.prev_loss = self.cur_loss

        # final save and vis step
        self.cur_step = num_iters
        self.save_checkpoint(out_dir)
        self.save_results(res_dir, seq_name)
        self.vis_result(res_dir, obs_data, vis)
        # self.save_meshes_all(res_dir, obs_data, seq_name)

    def optim_step(self, obs_data, i, writer=None):
        def closure():
            self.optim.zero_grad()
            #print('******************************', i)
            loss, stats_dict, preds = self.forward_pass(obs_data)
            #print('******************************', i)
            stats_dict["total"] = loss
            self.log_losses(move_to(detach_all(stats_dict), "cpu"))
            self.cur_loss = stats_dict["total"].detach().cpu().item()
            self.kkk = stats_dict
            self.ppp = preds
            loss_keys = stats_dict.keys()
            if "motion_prior" in loss_keys and "pose_prior" in loss_keys:
                self.cur_loss = (
                    self.cur_loss
                    - 0.04 * stats_dict["pose_prior"].detach().cpu().item()
                )
            print("optimizer print", loss, loss.shape, stats_dict, loss.requires_grad)
            loss.backward()
            # print('end of loss backward')
            return loss

        #print('!!!!!!!!!!!!!!!!!!!!!!!!!!', i)
        self.optim.step(closure)
        #print('!!!!!!!!!!!!!!!!!!!!!!!!!!', i)
        if writer is not None:
            self.record_current_losses(writer)


class RootOptimizer(StageOptimizer):
    name = "root_fit"
    stage = 0

    def __init__(
        self,
        model,
        all_loss_weights,
        use_chamfer=False,
        robust_loss_type="none",
        robust_tuning_const=4.6851,
        joints2d_sigma=100,
        **kwargs,
    ):
        param_names = ["trans", "root_orient"]
        super().__init__(self.name, model, param_names, **kwargs)

        self.loss = RootLoss(
            all_loss_weights[self.stage],
            ignore_op_joints=OP_IGNORE_JOINTS,
            joints2d_sigma=joints2d_sigma,
            use_chamfer=use_chamfer,
            robust_loss=robust_loss_type,
            robust_tuning_const=robust_tuning_const,
            faces=model.body_model.faces_tensor
        )

    def forward_pass(self, obs_data):
        """
        Takes in observed data, predicts the smpl parameters and returns loss
        """
        # raise ValueError
        # Use current params to go through SMPL and get joints3d, verts3d, points3d
        pred_data = self.model.pred_params_mano(obs_data["is_right"])
        pred_data["cameras"] = self.model.params.get_cameras()

        # compute data losses only
        vis_mask = obs_data["vis_mask"] >= 0
        # print('entering loss from RootOptimizer')
        loss, stats_dict = self.loss(obs_data, pred_data, vis_mask)
        return loss, stats_dict, pred_data


class SmoothOptimizer(StageOptimizer):
    name = "smooth_fit"
    stage = 1

    def __init__(
        self,
        model,
        all_loss_weights,
        use_chamfer=False,
        robust_loss_type="none",
        robust_tuning_const=4.6851,
        joints2d_sigma=100,
        **kwargs,
    ):
        param_names = ["trans", "root_orient", "betas", "latent_pose"]
        if model.opt_scale:
            param_names += ["world_scale"]
        if model.opt_cams:
            param_names += ["cam_f", "delta_cam_R"]

        super().__init__(self.name, model, param_names, **kwargs)

        self.loss = SMPLLoss(
            all_loss_weights[self.stage],
            ignore_op_joints=OP_IGNORE_JOINTS,
            joints2d_sigma=joints2d_sigma,
            use_chamfer=use_chamfer,
            robust_loss=robust_loss_type,
            robust_tuning_const=robust_tuning_const,
            faces=model.body_model.faces_tensor
        )

    def forward_pass(self, obs_data):
        """
        Takes in observed data, predicts the smpl parameters and returns loss
        """
        # raise ValueError
        # Use current params to go through SMPL and get joints3d, verts3d, points3d
        pred_data = self.model.pred_params_mano(obs_data["is_right"])
        pred_data["cameras"] = self.model.params.get_cameras()
        pred_data.update(self.model.params.get_vars())

        # camera predictions
        #         pts_ij, target, weight = self.model.params.get_reprojected_points()
        #         pred_data["bg2d_err"] = (
        #             (weight * (pts_ij - target) ** 2).sum(dim=(-1, -2)).mean()
        #         )
        pred_data["cam_R"], pred_data["cam_t"] = self.model.params.get_extrinsics()

        # compute data losses only
        vis_mask = obs_data["vis_mask"] >= 0
        loss, stats_dict = self.loss(obs_data, pred_data, self.model.seq_len, vis_mask)
        return loss, stats_dict, pred_data


# class MotionOptimizer(StageOptimizer):
#     name = "motion_fit"
#     stage = 2

#     def __init__(self, model, all_loss_weights, opt_cams=False, **kwargs):
#         self.opt_cams = opt_cams and model.opt_cams
#         param_names = [
#             "trans",
#             "root_orient",
#             "latent_pose",
#             "trans_vel",
#             "root_orient_vel",
#             "joints_vel",
#             "betas",
#             "latent_motion",
#             "floor_plane",
#         ]
#         if model.opt_scale:
#             param_names += ["world_scale"]

#         if self.opt_cams:
#             Logger.log(f"{self.name} OPTIMIZING CAMERAS")
#             param_names += ["delta_cam_R", "cam_f"]

#         super().__init__(self.name, model, param_names, **kwargs)
#         self.set_loss(model, all_loss_weights[self.stage], **kwargs)

#     def set_loss(
#         self,
#         model,
#         loss_weights,
#         use_chamfer=False,
#         robust_loss_type="none",
#         robust_tuning_const=4.6851,
#         joints2d_sigma=100,
#         **kwargs,
#     ):
#         self.loss = MotionLoss(
#             loss_weights,
#             init_motion_prior=model.init_motion_prior,
#             ignore_op_joints=OP_IGNORE_JOINTS,
#             joints2d_sigma=joints2d_sigma,
#             use_chamfer=use_chamfer,
#             robust_loss=robust_loss_type,
#             robust_tuning_const=robust_tuning_const,
#         )

#     def get_motion_scale(self):
#         return 1.0

#     def forward_pass(self, obs_data, num_steps=-1):
#         p = self.model.params
#         param_names = [
#             "betas",
#             "joints_vel",
#             "trans_vel",
#             "root_orient_vel",
#         ]
#         param_dict = p.get_vars(param_names)
#         # param_dict["floor_plane"] = p.floor_plane[p.floor_idcs]

#         preds, world_preds = self.model.rollout_mano_steps(num_steps)
#         preds.update(param_dict)
#         world_preds.update(param_dict)

#         # track mask is the length of the sequence that contains num_steps of each track
#         track_mask = world_preds.get("track_mask", None)
#         T = track_mask.shape[1] if track_mask is not None else num_steps
#         world_preds["cameras"] = p.get_cameras(np.arange(T))
#         if self.opt_cams:
#             cam_R, cam_t = p.get_extrinsics()
#             world_preds["cam_R"], world_preds["cam_t"] = cam_R[:T], cam_t[:T]

#         obs_data = slice_dict(obs_data, 0, T)
#         motion_scale = self.get_motion_scale()
#         loss, stats_dict = self.loss(
#             obs_data,
#             preds,
#             world_preds,
#             self.model.seq_len,
#             valid_mask=track_mask,
#             init_motion_scale=motion_scale,
#         )
#         return loss, stats_dict, (preds, world_preds)


# def slice_dict(d, start, end):
#     out = d.copy()
#     for k, v in d.items():
#         if not isinstance(v, torch.Tensor) or v.ndim < 3 or v.shape[1] < end:
#             continue
#         out[k] = v[:, start:end]
#     return out


# class MotionOptimizerChunks(MotionOptimizer):
#     """
#     Incrementally optimize sequence in chunks (with all variables)
#     :param chunk_size (int) number of frames per chunk
#     :param init_steps (int) number of optimization steps to spend on first chunk
#     :param chunk_steps (int) number of opt steps to spend on subsequent chunks
#     """

#     name = "motion_chunks"
#     stage = 2

#     def __init__(self, *args, chunk_size=20, init_steps=20, chunk_steps=10, **kwargs):
#         self.chunk_size = chunk_size
#         self.init_steps = init_steps
#         self.chunk_steps = chunk_steps
#         super().__init__(*args, **kwargs)

#     @property
#     def num_iters(self):
#         num_chunks = int(np.ceil(self.model.seq_len / self.chunk_size)) + 1
#         return self.init_steps + self.chunk_steps * num_chunks

#     @property
#     def end_idx(self):
#         if self.cur_step < self.init_steps and self.add_chunk == 0:
#             return self.chunk_size
#         chunk_idx = max(
#             (self.cur_step - self.init_steps) // self.chunk_steps + 1, self.add_chunk
#         )
#         if int(np.ceil(self.model.seq_len / self.chunk_size)) == self.add_chunk:
#             self.reached_max = True
#         return min(self.chunk_size * (chunk_idx + 1), self.model.seq_len)

#     @property
#     def start_idx(self):
#         return 0

#     def get_motion_scale(self):
#         num_frames = self.end_idx - self.start_idx
#         return max(1.0, float(self.model.seq_len) / num_frames)

#     def forward_pass(self, obs_data):
#         return super().forward_pass(obs_data, num_steps=self.end_idx)

#     def vis_result(self, res_dir, obs_data, vis=None, **kwargs):
#         print("start, end", self.start_idx, self.end_idx)
#         return super().vis_result(res_dir, obs_data, vis=vis, num_steps=self.end_idx)


# class MotionOptimizerFreeze(MotionOptimizer):
#     """
#     Only optimizing a subset of parameters
#     """

#     name = "motion_freeze"
#     stage = 2

#     def __init__(self, model, all_loss_weights, no_contacts=True, **kwargs):
#         loss_weights = all_loss_weights[self.stage].copy()

#         if no_contacts:
#             loss_weights["contact_height"] = 0.0
#             loss_weights["contact_vel"] = 0.0

#         param_names = ["latent_motion", "betas", "trans", "floor_plane"]
#         if model.optim_scale:
#             param_names += ["world_scale"]

#         StageOptimizer.__init__(self, self.name, model, param_names, **kwargs)
#         self.set_loss(model, loss_weights, **kwargs)