import argparse
from typing import Dict, List
import gin
from pypsdr.renderer import Model
import psdr_cpu
import dataclasses
from dataclasses import field
from pypsdr.validate import *
from pypsdr.utils.io import *
from pypsdr.loss import compute_render_loss, uniform_laplacian
from pypsdr.optimize import Camera
from pypsdr.optimizer import LGDescent, sparse_eye
from pypsdr.renderer import Model, Render
from pypsdr.common import gen_cameras
from pypsdr.largesteps.parameterize import from_differential, to_differential
from pypsdr.largesteps.geometry import compute_matrix
from pypsdr.largesteps.optimize import AdamUniform
from random import randrange, seed

import scipy
import scipy.sparse.linalg
import datetime


@gin.configurable
# generate the constructor
@dataclasses.dataclass
class TrainRunner:
    scene_init_file: str = "./kitty/ini.xml"
    scene_target_file: str = "./kitty/tar.xml"
    niter: int = 500
    options: psdr_cpu.RenderOptions = None
    out_dir: str = "./output"          # output directory
    # directory that contains target images
    target_dir: str = "./kitty/output/kitty/target"
    test_indices: List = field(
        default_factory=lambda: [0, 14, 32, 57])
    img_threshold: float = 2.
    batch_size: int = 10
    print_size: int = 10
    train_indices: np.array = None
    paramters: List = field(default_factory=lambda: [])

    camera_file: str = "./kitty/cam_pos.txt"
    integrator: psdr_cpu.Integrator = psdr_cpu.Direct()
    parameters: Dict = field(default_factory=lambda: {})

    def __post_init__(self):
        mkdir(self.out_dir)
        mkdir(os.path.join(self.out_dir, "iter"))
        self.scene_init = Scene(self.scene_init_file)
        self.scene_tar = Scene(self.scene_target_file)
        self.width = self.scene_init.camera.width
        self.height = self.scene_init.camera.height
        self.fov = self.scene_init.camera.fov
        self.train_images = np.array(read_images(self.target_dir))
        self.test_images = self.train_images[self.test_indices]
        self.test_image = np.concatenate(self.test_images, axis=1)
        imwrite(self.test_image,
                os.path.join(self.out_dir, "test.png"))

        bitmap = self.scene_tar.bsdfs[2].reflectance
        bitmap = np.array(bitmap.m_data).reshape(
            bitmap.m_res[1], bitmap.m_res[0], 3)
        imwrite(bitmap, os.path.join(
                self.out_dir, "text.exr"))

        self.train_images = torch.from_numpy(self.train_images)
        self.test_images = torch.from_numpy(self.test_images)

        if self.train_indices:
            self.train_indices = np.array(self.train_indices)
            assert(self.train_indices[self.train_indices >= len(
                self.train_images)].size == 0)
        else:
            self.train_indices = np.arange(len(self.train_images))

        camera_pos = np.loadtxt(self.camera_file)
        camera_info = Camera(target=gin.REQUIRED)

        self.cameras = gen_cameras(positions=camera_pos,
                                   target=camera_info.target,
                                   up=camera_info.up,
                                   fov=camera_info.fov,
                                   resolution=[self.width, self.height])
        self.sceneAD = psdr_cpu.SceneAD(self.scene_init)
        context = {
            "integrator": self.integrator,
            "options": self.options,
            "sceneAD": self.sceneAD,
            "cameras": self.cameras
        }
        self.model = Model(self.sceneAD, self.parameters, context=context)
        # psdr_cpu.set_verbose(True)

    def run(self):
        # setup optimizers
        print('Starting optimization')
        seed(1)
        np.random.seed(0)
        # generate mini-batches
        idx_shuffled = np.array([], dtype=np.int32)
        for i in range(self.batch_size * self.niter // len(self.train_indices) + 1):
            shuffled = self.train_indices.copy()
            np.random.shuffle(shuffled)
            idx_shuffled = np.concatenate((idx_shuffled, shuffled))

        error = []
        for iter in range(self.niter):
            now = datetime.datetime.now()
            self.model.zero_grad()
            loss = 0.0
            indices = idx_shuffled[iter *
                                   self.batch_size: (iter + 1) * self.batch_size]
            for j in range(self.batch_size):
                sensor_id = indices[j]
                # render
                img = self.model.render(context={"sensor_id": sensor_id,
                                                 "options": RenderOptions(iter,
                                                                          gin.REQUIRED,
                                                                          gin.REQUIRED,
                                                                          gin.REQUIRED)})
                tar = self.train_images[sensor_id]
                # clamping
                img = img.clamp(0., self.img_threshold)
                tar = tar.clamp(0., self.img_threshold)
                img_loss = compute_render_loss(
                    img, tar, 1.0)
                print("image_loss: ", img_loss.item())
                loss = img_loss.item()
                img_loss.backward()
                print("Rendering camera:", sensor_id,
                      "loss:", img_loss.item(), end='\r')

            # step
            self.model.step()

            # TODO update scene with new parameters

            # print stat
            # print("grad: ", np.abs(u.grad.cpu()).sum().item())
            end = datetime.datetime.now() - now
            print("rendered camera:", indices)
            print('Iteration:', iter,
                  'Loss:', loss,
                  'Total time:', end.seconds + end.microseconds / 1000000.0)
            error.append(loss)
            np.savetxt(os.path.join(self.out_dir, "loss.log"), error)

            # TODO save parameters

            # TODO write files
            # self.model.save()
            bitmap = self.sceneAD.val.bsdfs[2].reflectance
            bitmap = np.array(bitmap.m_data).reshape(
                bitmap.m_res[1], bitmap.m_res[0], 3)
            imwrite(bitmap, os.path.join(
                self.out_dir, "iter", "text_%04d.exr" % iter))

            if iter % self.print_size == 0:
                test_images = []
                for test_id in self.test_indices:
                    print("Testing camera:", test_id, end='\r')
                    test_img = self.model.render(
                        context={"sensor_id": test_id,
                                 "options": RenderOptions(iter, gin.REQUIRED, 128, 0, 0)}).detach()
                    test_images.append(test_img)
                test_image = np.concatenate(test_images, axis=1)
                imwrite(np.concatenate([self.test_image, test_image], axis=0),
                        os.path.join(self.out_dir, "iter", "iter_%d.png" % iter))


if __name__ == "__main__":
    default_config = './texture2.conf'
    parser = argparse.ArgumentParser(
        description='Script for generating validation results')
    parser.add_argument('config_file', metavar='config_file',
                        type=str, nargs='?', default=default_config, help='config file')
    args, unknown = parser.parse_known_args()
    # Dependency injection: Arguments are injected into the function from the gin config file.
    gin.parse_config_file(args.config_file, skip_unknown=True)
    if os.path.dirname(args.config_file) != '':
        os.chdir(os.path.dirname(args.config_file))
    train_runner = TrainRunner()
    train_runner.run()
