from custom_types import *
import options as options
from process_data import mesh_utils
import models.factory as factory
from models.single_mesh_models import SingleMeshGenerator
from models.mesh_handler import MeshHandler, MeshInference, load_template_mesh


class DGTS:

    def __init__(self, opt: Union[options.Options, options.TrainOption], device: D):
        self.opt = opt.load()
        self.generator: SingleMeshGenerator = factory.model_lc(opt, SingleMeshGenerator, device=device)
        self.level = 0
        self.device = device

    def get_random_z(self, num_randoms: int) -> T:
        if self.opt.noise_before:
            return torch.randn(num_randoms, 3, device=self.device) * self.opt.noise_amplitude
        return torch.randn(1, self.generator.opt.in_nf, num_randoms, device=self.device) * self.opt.noise_amplitude

    def get_z_by_level(self, base_mesh: MeshHandler, level: int) -> T:
        num_faces = len(base_mesh) * 4 ** level
        if self.opt.noise_before:
            num_faces = (num_faces - len(base_mesh)) // 2 + base_mesh.vs.shape[0]
        return self.get_random_z(num_faces)

    def get_z_sequence(self, base_mesh: MeshHandler, max_level: int) -> TS:
        return [self.get_z_by_level(base_mesh, level) for level in range(max_level + 1)]

    def __len__(self):
        return len(self.generator.levels)


class Mesh2Mesh(DGTS):

    def __init__(self, opt: options.Options, device: D):
        super(Mesh2Mesh, self).__init__(opt, device)
        self.generator.eval()

    def trim(self, start: int, end: int) -> Tuple[int, int]:
        if end < 0 or end > len(self) - 1:
                end = len(self) - 1
        if start > end:
            start = end
        return start, end

    def get_z_sequence(self, base_mesh: MeshHandler, max_level: int) -> factory.Noise:
        return factory.Noise(data=super(Mesh2Mesh, self).get_z_sequence(base_mesh, max_level))

    def growing(self, mesh: MeshInference, start: int, end: int, num_frames: int, zero_places: NoiseT = ()) ->  factory.Noise:

        export_name = f'{self.opt.cp_folder}/inference/{mesh.mesh_name}/scene00'

        start, end = self.trim(start, end)
        if len(zero_places) == 1:
            zero_places = zero_places * (end - start + 1)
        z = self.get_z_sequence(mesh, end - start)
        for i in range(min(len(z), len(zero_places))):
            if zero_places[i]:
                z[i] = 0
        deltas = self.generator.grow_forward(mesh.copy(), z, end, start)
        mesh.export(f'{export_name}/{0:02d}')
        for i in range(end - start + 1):
            base_vs, cur_delta = mesh.vs.clone(), deltas[i]
            for j in range(num_frames):
                mesh.vs = base_vs + cur_delta * (j + 1) / num_frames
                mesh.export(f'{export_name}/{(num_frames * i + j + 1):02d}')
                print(f'done: {num_frames * i + j + 1}/{num_frames * (end - start + 1)}')
            if i < end - start:
                mesh.upsample()
        return z

    def animate(self, mesh: MeshInference, start: int, end: int, num_scene: int, num_frames: Tuple[int, int], zero_places: NoiseT = ()):
        export_name = f'{self.opt.cp_folder}/inference/{mesh.mesh_name}/scene01'
        start, end = self.trim(start, end)
        if len(zero_places) == 1:
            zero_places = zero_places * (end - start + 1)
        z_a = self.growing(mesh.copy(), start, end, num_frames[0], zero_places)
        z_b = self.get_z_sequence(mesh, end - start).to(self.device)
        for i in range(min(len(z_a), len(zero_places))):
            if zero_places[i]:
                z_b[i] = 0
        z_start = z_a
        num_frames = num_frames[1]
        for s in range(num_scene):
            for i in range(num_frames):
                alpha = (i + 1) / float(num_frames)
                z = z_a * (1 - alpha) + z_b * alpha
                m = mesh.copy()
                out = self.generator(m, z, end, start, upsample=True)
                out.export(f'{export_name}/{s * num_frames + i:02d}')
                print(f'frame {s * num_frames + i + 1} / {num_scene * num_frames}...')

            z_a = z_b
            if s == num_scene - 2:
                z_b = z_start
            else:
                z_b = self.get_z_sequence(mesh, end - start).to(self.device)

    def __call__(self, mesh: Union[str, MeshHandler, T_Mesh], start: int, end: int, zero_places: NoiseT = 0) -> MeshHandler:
        MeshHandler.reset()
        start, end = self.trim(start, end)
        if type(zero_places) is int:
            zero_places = [zero_places]
        if len(zero_places) == 1:
            zero_places = zero_places * (end - start + 1)
        if mesh is None:
            mesh = MeshHandler(mesh_utils.load_real_mesh(self.opt.template_name, start),self.opt, 0).to(self.device)
        elif type(mesh) is not MeshHandler:
            mesh = MeshHandler(mesh, self.opt, 0).to(self.device)
        z = self.get_z_sequence(mesh, end - start)
        for i in range(min(len(z), len(zero_places))):
            if zero_places[i]:
                z[i] = 0
        remeshed = self.generator.forward(mesh, z, end, start, upsample=True)
        return remeshed


class MeshGen(DGTS):

    def __init__(self, opt: options.Options, device: D):
        super(MeshGen, self).__init__(opt, device)
        self.generator.eval()
        template_name, template = load_template_mesh(opt, opt.start_level)
        self.template = MeshInference(template_name, template, self.opt, self.opt.start_level).to(self.device)
        self.reconstruction_z = factory.NoiseMem(opt).load().to(device)

    def compose_z(self, start_level) -> factory.Noise:
        random_noise = self.get_z_sequence(self.template, len(self) - 1)
        noise = self.reconstruction_z[: start_level] + random_noise[start_level:]
        return noise

    def generate_seq(self, num_seqs: int):
        for seq in range(num_seqs):
            z = self.compose_z(0)
            self.generator.inference_forward(self.template.copy(), z, len(self) - 1, 0,
                                             f'{opt_.cp_folder}/inference/gen/{self.opt.mesh_name}_{seq}',
                                             upsample=True)

    def generate_all(self, num_samples: int):
        for i in range(len(self.generator.levels)):
            for j in range(num_samples):
                out_mesh = self(i)
                out_mesh.export(f'{self.opt.cp_folder}/inference/gen/{self.opt.mesh_name}_{i}_{j:02d}')
                print(f'gen {self.opt.mesh_name} {i * num_samples + j +1:02d} / {len(self.generator.levels) * num_samples}')

    def __call__(self, start_level: int):
        with torch.no_grad():
            if start_level < 0:
                start_level = len(self)
            start_level = min(len(self), start_level)
            z = self.compose_z(start_level)
            return self.generator(self.template.copy(), z, len(self) - 1)


if __name__ == '__main__':
    opt_ = options.Options()
    opt_.parse_cmdline()
    device = CPU
    with_noise = False
    if opt_.gen_mode == 'generate':
        mg = MeshGen(opt_, device)
        mg.generate_all(opt_.num_gen_samples)
    elif opt_.gen_mode == 'animate':
        m2m = Mesh2Mesh(opt_, device)
        in_mesh = MeshInference(opt_.target, mesh_utils.load_real_mesh(opt_.target, 0, True), opt_, 0).to(device)
        m2m.animate(in_mesh, opt_.gen_levels[0], opt_.gen_levels[1], 0, (12, 17), zero_places=(0, 0, 1, 1, 1))