#!/usr/bin/env python3
# Copyright 2004-present Facebook. All Rights Reserved.

import json
import os
import torch

model_params_subdir = "ModelParameters"
optimizer_params_subdir = "OptimizerParameters"
latent_codes_subdir = "LatentCodes"
logs_filename = "Logs.pth"
reconstructions_subdir = "Reconstructions"
reconstruction_meshes_subdir = "Meshes"
reconstruction_codes_subdir = "Codes"
optimizations_subdir = "Optimizations"
optimizations_meshes_subdir = "Meshes"
optimizations_codes_subdir = "Codes"
specifications_filename = "specs.json"
data_source_map_filename = ".datasources.json"
evaluation_subdir = "Evaluation"
sdf_samples_subdir = "SdfSamples"
renders_subdir = "Renders"
surface_samples_subdir = "SurfaceSamples"
normalization_param_subdir = "NormalizationParameters"
training_meshes_subdir = "TrainingMeshes"


def load_experiment_specifications(experiment_directory):

    filename = os.path.join(experiment_directory, specifications_filename)

    if not os.path.isfile(filename):
        raise Exception(
            "The experiment directory ({}) does not include specifications file "
            + '"specs.json"'.format(experiment_directory)
        )

    return json.load(open(filename))


def load_model_parameters(experiment_directory, checkpoint, decoder):

    filename = os.path.join(
        experiment_directory, model_params_subdir, checkpoint + ".pth"
    )

    if not os.path.isfile(filename):
        raise Exception('model state dict "{}" does not exist'.format(filename))

    data = torch.load(filename)

    decoder.load_state_dict(data["model_state_dict"])

    return data["epoch"]


def build_decoder(experiment_directory, experiment_specs):

    arch = __import__(
        "networks." + experiment_specs["NetworkArch"], fromlist=["Decoder"]
    )

    latent_size = experiment_specs["CodeLength"]

    decoder = arch.Decoder(latent_size, **experiment_specs["NetworkSpecs"]).cuda()

    return decoder


def load_decoder(
    experiment_directory, experiment_specs, checkpoint, data_parallel=True
):

    decoder = build_decoder(experiment_directory, experiment_specs)

    if data_parallel:
        decoder = torch.nn.DataParallel(decoder)

    epoch = load_model_parameters(experiment_directory, checkpoint, decoder)

    return (decoder, epoch)


def load_latent_vectors(experiment_directory, checkpoint):

    filename = os.path.join(
        experiment_directory, latent_codes_subdir, checkpoint + ".pth"
    )

    if not os.path.isfile(filename):
        raise Exception(
            "The experiment directory ({}) does not include a latent code file"
            + " for checkpoint '{}'".format(experiment_directory, checkpoint)
        )

    data = torch.load(filename)

    if isinstance(data["latent_codes"], torch.Tensor):

        num_vecs = data["latent_codes"].size()[0]

        lat_vecs = []
        for i in range(num_vecs):
            lat_vecs.append(data["latent_codes"][i].cuda())

        return lat_vecs

    else:

        num_embeddings, embedding_dim = data["latent_codes"]["weight"].shape

        lat_vecs = torch.nn.Embedding(num_embeddings, embedding_dim)

        lat_vecs.load_state_dict(data["latent_codes"])

        return lat_vecs.weight.data.detach()


def get_data_source_map_filename(data_dir):
    return os.path.join(data_dir, data_source_map_filename)


def get_reconstructed_mesh_filename(
    experiment_dir, epoch, dataset, class_name, instance_name
):

    return os.path.join(
        experiment_dir,
        reconstructions_subdir,
        str(epoch),
        reconstruction_meshes_subdir,
        dataset,
        class_name,
        instance_name + ".ply",
    )


def get_reconstructed_code_filename(
    experiment_dir, epoch, dataset, class_name, instance_name
):

    return os.path.join(
        experiment_dir,
        reconstructions_subdir,
        str(epoch),
        reconstruction_codes_subdir,
        dataset,
        class_name,
        instance_name + ".pth",
    )


def get_evaluation_dir(experiment_dir, checkpoint, create_if_nonexistent=False):

    dir = os.path.join(experiment_dir, evaluation_subdir, checkpoint)

    if create_if_nonexistent and not os.path.isdir(dir):
        os.makedirs(dir)

    return dir


def get_model_params_dir(experiment_dir, create_if_nonexistent=False):

    dir = os.path.join(experiment_dir, model_params_subdir)

    if create_if_nonexistent and not os.path.isdir(dir):
        os.makedirs(dir)

    return dir


def get_optimizer_params_dir(experiment_dir, create_if_nonexistent=False):

    dir = os.path.join(experiment_dir, optimizer_params_subdir)

    if create_if_nonexistent and not os.path.isdir(dir):
        os.makedirs(dir)

    return dir


def get_latent_codes_dir(experiment_dir, create_if_nonexistent=False):

    dir = os.path.join(experiment_dir, latent_codes_subdir)

    if create_if_nonexistent and not os.path.isdir(dir):
        os.makedirs(dir)

    return dir


def get_normalization_params_filename(
    data_dir, dataset_name, class_name, instance_name
):
    return os.path.join(
        data_dir,
        normalization_param_subdir,
        dataset_name,
        class_name,
        instance_name + ".npz",
    )

[((30, 15, 30, 74), 'os.path.join', 'os.path.join', ({(30, 28, 30, 48): 'experiment_directory', (30, 50, 30, 73): 'specifications_filename'}, {}), '(experiment_directory, specifications_filename)', False, 'import os\n'), ((43, 15, 45, 5), 'os.path.join', 'os.path.join', ({(44, 8, 44, 28): 'experiment_directory', (44, 30, 44, 49): 'model_params_subdir', (44, 51, 44, 70): "checkpoint + '.pth'"}, {}), "(experiment_directory, model_params_subdir, checkpoint + '.pth')", False, 'import os\n'), ((50, 11, 50, 31), 'torch.load', 'torch.load', ({(50, 22, 50, 30): 'filename'}, {}), '(filename)', False, 'import torch\n'), ((86, 15, 88, 5), 'os.path.join', 'os.path.join', ({(87, 8, 87, 28): 'experiment_directory', (87, 30, 87, 49): 'latent_codes_subdir', (87, 51, 87, 70): "checkpoint + '.pth'"}, {}), "(experiment_directory, latent_codes_subdir, checkpoint + '.pth')", False, 'import os\n'), ((96, 11, 96, 31), 'torch.load', 'torch.load', ({(96, 22, 96, 30): 'filename'}, {}), '(filename)', False, 'import torch\n'), ((120, 11, 120, 59), 'os.path.join', 'os.path.join', ({(120, 24, 120, 32): 'data_dir', (120, 34, 120, 58): 'data_source_map_filename'}, {}), '(data_dir, data_source_map_filename)', False, 'import os\n'), ((155, 10, 155, 69), 'os.path.join', 'os.path.join', ({(155, 23, 155, 37): 'experiment_dir', (155, 39, 155, 56): 'evaluation_subdir', (155, 58, 155, 68): 'checkpoint'}, {}), '(experiment_dir, evaluation_subdir, checkpoint)', False, 'import os\n'), ((165, 10, 165, 59), 'os.path.join', 'os.path.join', ({(165, 23, 165, 37): 'experiment_dir', (165, 39, 165, 58): 'model_params_subdir'}, {}), '(experiment_dir, model_params_subdir)', False, 'import os\n'), ((175, 10, 175, 63), 'os.path.join', 'os.path.join', ({(175, 23, 175, 37): 'experiment_dir', (175, 39, 175, 62): 'optimizer_params_subdir'}, {}), '(experiment_dir, optimizer_params_subdir)', False, 'import os\n'), ((185, 10, 185, 59), 'os.path.join', 'os.path.join', ({(185, 23, 185, 37): 'experiment_dir', (185, 39, 185, 58): 'latent_codes_subdir'}, {}), '(experiment_dir, latent_codes_subdir)', False, 'import os\n'), ((196, 11, 202, 5), 'os.path.join', 'os.path.join', ({(197, 8, 197, 16): 'data_dir', (198, 8, 198, 34): 'normalization_param_subdir', (199, 8, 199, 20): 'dataset_name', (200, 8, 200, 18): 'class_name', (201, 8, 201, 30): "(instance_name + '.npz')"}, {}), "(data_dir, normalization_param_subdir, dataset_name, class_name,\n    instance_name + '.npz')", False, 'import os\n'), ((32, 11, 32, 35), 'os.path.isfile', 'os.path.isfile', ({(32, 26, 32, 34): 'filename'}, {}), '(filename)', False, 'import os\n'), ((47, 11, 47, 35), 'os.path.isfile', 'os.path.isfile', ({(47, 26, 47, 34): 'filename'}, {}), '(filename)', False, 'import os\n'), ((77, 18, 77, 48), 'torch.nn.DataParallel', 'torch.nn.DataParallel', ({(77, 40, 77, 47): 'decoder'}, {}), '(decoder)', False, 'import torch\n'), ((90, 11, 90, 35), 'os.path.isfile', 'os.path.isfile', ({(90, 26, 90, 34): 'filename'}, {}), '(filename)', False, 'import os\n'), ((112, 19, 112, 68), 'torch.nn.Embedding', 'torch.nn.Embedding', ({(112, 38, 112, 52): 'num_embeddings', (112, 54, 112, 67): 'embedding_dim'}, {}), '(num_embeddings, embedding_dim)', False, 'import torch\n'), ((158, 8, 158, 24), 'os.makedirs', 'os.makedirs', ({(158, 20, 158, 23): 'dir'}, {}), '(dir)', False, 'import os\n'), ((168, 8, 168, 24), 'os.makedirs', 'os.makedirs', ({(168, 20, 168, 23): 'dir'}, {}), '(dir)', False, 'import os\n'), ((178, 8, 178, 24), 'os.makedirs', 'os.makedirs', ({(178, 20, 178, 23): 'dir'}, {}), '(dir)', False, 'import os\n'), ((188, 8, 188, 24), 'os.makedirs', 'os.makedirs', ({(188, 20, 188, 23): 'dir'}, {}), '(dir)', False, 'import os\n'), ((157, 37, 157, 55), 'os.path.isdir', 'os.path.isdir', ({(157, 51, 157, 54): 'dir'}, {}), '(dir)', False, 'import os\n'), ((167, 37, 167, 55), 'os.path.isdir', 'os.path.isdir', ({(167, 51, 167, 54): 'dir'}, {}), '(dir)', False, 'import os\n'), ((177, 37, 177, 55), 'os.path.isdir', 'os.path.isdir', ({(177, 51, 177, 54): 'dir'}, {}), '(dir)', False, 'import os\n'), ((187, 37, 187, 55), 'os.path.isdir', 'os.path.isdir', ({(187, 51, 187, 54): 'dir'}, {}), '(dir)', False, 'import os\n')]