#!/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",
    )

[((898, 957), 'os.path.join', 'os.path.join', (['experiment_directory', 'specifications_filename'], {}), '(experiment_directory, specifications_filename)\n', (910, 957), False, 'import os\n'), ((1297, 1373), 'os.path.join', 'os.path.join', (['experiment_directory', 'model_params_subdir', "(checkpoint + '.pth')"], {}), "(experiment_directory, model_params_subdir, checkpoint + '.pth')\n", (1309, 1373), False, 'import os\n'), ((1519, 1539), 'torch.load', 'torch.load', (['filename'], {}), '(filename)\n', (1529, 1539), False, 'import torch\n'), ((2365, 2441), 'os.path.join', 'os.path.join', (['experiment_directory', 'latent_codes_subdir', "(checkpoint + '.pth')"], {}), "(experiment_directory, latent_codes_subdir, checkpoint + '.pth')\n", (2377, 2441), False, 'import os\n'), ((2699, 2719), 'torch.load', 'torch.load', (['filename'], {}), '(filename)\n', (2709, 2719), False, 'import torch\n'), ((3287, 3335), 'os.path.join', 'os.path.join', (['data_dir', 'data_source_map_filename'], {}), '(data_dir, data_source_map_filename)\n', (3299, 3335), False, 'import os\n'), ((4067, 4126), 'os.path.join', 'os.path.join', (['experiment_dir', 'evaluation_subdir', 'checkpoint'], {}), '(experiment_dir, evaluation_subdir, checkpoint)\n', (4079, 4126), False, 'import os\n'), ((4310, 4359), 'os.path.join', 'os.path.join', (['experiment_dir', 'model_params_subdir'], {}), '(experiment_dir, model_params_subdir)\n', (4322, 4359), False, 'import os\n'), ((4547, 4600), 'os.path.join', 'os.path.join', (['experiment_dir', 'optimizer_params_subdir'], {}), '(experiment_dir, optimizer_params_subdir)\n', (4559, 4600), False, 'import os\n'), ((4784, 4833), 'os.path.join', 'os.path.join', (['experiment_dir', 'latent_codes_subdir'], {}), '(experiment_dir, latent_codes_subdir)\n', (4796, 4833), False, 'import os\n'), ((5042, 5146), 'os.path.join', 'os.path.join', (['data_dir', 'normalization_param_subdir', 'dataset_name', 'class_name', "(instance_name + '.npz')"], {}), "(data_dir, normalization_param_subdir, dataset_name, class_name,\n    instance_name + '.npz')\n", (5054, 5146), False, 'import os\n'), ((970, 994), 'os.path.isfile', 'os.path.isfile', (['filename'], {}), '(filename)\n', (984, 994), False, 'import os\n'), ((1400, 1424), 'os.path.isfile', 'os.path.isfile', (['filename'], {}), '(filename)\n', (1414, 1424), False, 'import os\n'), ((2150, 2180), 'torch.nn.DataParallel', 'torch.nn.DataParallel', (['decoder'], {}), '(decoder)\n', (2171, 2180), False, 'import torch\n'), ((2468, 2492), 'os.path.isfile', 'os.path.isfile', (['filename'], {}), '(filename)\n', (2482, 2492), False, 'import os\n'), ((3078, 3127), 'torch.nn.Embedding', 'torch.nn.Embedding', (['num_embeddings', 'embedding_dim'], {}), '(num_embeddings, embedding_dim)\n', (3096, 3127), False, 'import torch\n'), ((4193, 4209), 'os.makedirs', 'os.makedirs', (['dir'], {}), '(dir)\n', (4204, 4209), False, 'import os\n'), ((4426, 4442), 'os.makedirs', 'os.makedirs', (['dir'], {}), '(dir)\n', (4437, 4442), False, 'import os\n'), ((4667, 4683), 'os.makedirs', 'os.makedirs', (['dir'], {}), '(dir)\n', (4678, 4683), False, 'import os\n'), ((4900, 4916), 'os.makedirs', 'os.makedirs', (['dir'], {}), '(dir)\n', (4911, 4916), False, 'import os\n'), ((4165, 4183), 'os.path.isdir', 'os.path.isdir', (['dir'], {}), '(dir)\n', (4178, 4183), False, 'import os\n'), ((4398, 4416), 'os.path.isdir', 'os.path.isdir', (['dir'], {}), '(dir)\n', (4411, 4416), False, 'import os\n'), ((4639, 4657), 'os.path.isdir', 'os.path.isdir', (['dir'], {}), '(dir)\n', (4652, 4657), False, 'import os\n'), ((4872, 4890), 'os.path.isdir', 'os.path.isdir', (['dir'], {}), '(dir)\n', (4885, 4890), False, 'import os\n')]