#Script which takes all autolabels generated
#under a given sequence root, and builds TFrecord files
#(for depth image classifier training)
#This does no shuffling -- TODO: handle that with a separate utility
import sys
import pickle
import ColorFilters
import numpy as np
from RootDirector import *
from DepthTrainingTFWriter import *
import StandardBody
from scipy.spatial import cKDTree
from multiprocessing import Process, Pool
import traceback
import gc

NUM_PROCS = 5

def batcher(seqNameList):
    result = []
    mostRecent = []
    for seqName in seqNameList:
        mostRecent.append(seqName)
        if (len(mostRecent) >= NUM_PROCS):
            result.append(mostRecent)
            mostRecent = []
    result.append(mostRecent)
    return result

#For the alternative template thing here, get those indices on the
#template body which correspond to the torso

coloredTemplateFile = "ColoredTemplate.pickle"
coloredBody = pickle.load(open(coloredTemplateFile, "rb"))
colorArray = np.copy(np.asarray(coloredBody.colors))

torsoMask = ColorFilters.pixelSpaceVectorMask(ColorFilters.maskTorso, colorArray[:, 0:3])
torsoMask = np.logical_not(torsoMask)
torsoInds = np.where(torsoMask)

def alternativeGetTemplateIndexImage(pointCloud, labels):
    try:
        #Does all of the stuff you'd expect here, but first tries
        #to align the pointCloud and labels with a translation

        #First, find the points on the torso in the point cloud (if any)
        #If there aren't any, then we just defer to the default implementation
        torsoCloud = pointCloud.copy()
        torsoCloud.applyColorFilter(ColorFilters.maskTorso, negated=True)

        torsoCloudArray = np.asarray(torsoCloud.getPoints())
        torsoLabelArray = labels[torsoInds]

        if (torsoCloudArray.size < 6):
            return getTemplateIndexImage(pointCloud, labels)

        quantile_q = .1

        #Get the qth quantile of both the point cloud and labels z, and make them align
        #by adjusting the labels as needed
        labelZQuantile = np.quantile(torsoLabelArray[:, 2], quantile_q)

        pointCloudZQuantile = np.quantile(torsoCloudArray[:, 2], quantile_q)

        z_adjust = pointCloudZQuantile - labelZQuantile

        print z_adjust

        labels = np.copy(labels) + np.array([[0.0, 0.0, z_adjust]])

        return getTemplateIndexImage(pointCloud, labels)

    except:
        raise Exception("".join(traceback.format_exception(*sys.exc_info())))

def advancedGetTemplateIndexImage(pointCloud, labels):

    try:
        #getTemplateIndexImage did things relatively simple, in
        #the sense that there, it was just a one-time assignment
        #of the closest distorted template point from points
        #in the point cloud

        #Here, we'll kinda-sorta do that, but iteratively
        #reassign points in the point cloud array to the average
        #of nearby neighbors (in template-space), weighted by probabilities given
        #by the original distance from the each point to the label cloud
        #We'll then snap those positions to positions on the template manifold,
        #take those as the new coordinates, and do it again.
        #The result is hopefully smoother and more coherent than the original,
        #and hopefully majority-takes-all for mislabelings

        labelKdTree = cKDTree(labels)

        cloudPointArray = np.asarray(pointCloud.getPoints(), dtype=np.float32)

        K = 8
        if (cloudPointArray.shape[0] < K):
            #Special case: not enough points to do the averaging!
            return getTemplateIndexImage(pointCloud, labels)

        #Get the indices of a K-neighborhood of every point in the input cloud
        cloudKdTree = cKDTree(cloudPointArray)
        _, cloudKNeighborhoods = cloudKdTree.query(cloudPointArray, K)

        labelDists, labelInds = labelKdTree.query(cloudPointArray)

        #Assign probability weights based on a decreasing exponential
        #of distance
        #Distance (mm) which constitutes a weight reduction to 1/e of what it was
        naturalDist = 50.0
        probWeights = np.exp(-labelDists / naturalDist)

        NUM_AVERAGING_STEPS = 10

        for _ in range(NUM_AVERAGING_STEPS):
            #Get corresponding points on the standard template
            templatePoints = StandardBody.pointArray[labelInds]
            #Space to store the results
            newTemplatePoints = np.copy(templatePoints)
            #TODO: Can this loop be sped up?
            for i in range(cloudPointArray.shape[0]):
                neighborInds = cloudKNeighborhoods[i]
                neighborTemplatePoints = templatePoints[neighborInds]
                neighborProbWeights = probWeights[neighborInds]
                #Okay, now compute a weighted average to determine the new template point pos
                newTemplatePoint = np.average(neighborTemplatePoints, axis=0, weights=neighborProbWeights)
                newTemplatePoints[i] = newTemplatePoint

            #Great, now using the new template points, find the closest labelInds by searching
            #the standard template kd Tree
            _, labelInds = StandardBody.standardKdTree.query(newTemplatePoints)

        cloudDepthIndexArray = np.asarray(pointCloud.getIndices(), dtype=np.int32)

        cloudDepthYindices = cloudDepthIndexArray[:, 0].flatten()
        cloudDepthXindices = cloudDepthIndexArray[:, 1].flatten()

        #Convert label indices to uint16 for output
        labelInds = labelInds.astype(np.uint16)

        result = np.full((424, 512), StandardBody.pointArray.shape[0], dtype=np.uint16)

        result[cloudDepthYindices, cloudDepthXindices] = labelInds

        #TODO: Void labels that aren't close to their constituent parts,
        #and void or correct whole frames which get the backside /frontside thing wrong

        return result
    except:
        raise Exception("".join(traceback.format_exception(*sys.exc_info())))

#Given a RGBPointCloud and corresponding standard body labels,
#return a 424x512 uint16 numpy array of the original depth frame's
#closest matched template points
def getTemplateIndexImage(pointCloud, labels):
    #First thing's first, initialize the result array to "not present"
    #values (taken as StandardBody.pointArray.shape[0])

    result = np.full((424, 512), StandardBody.pointArray.shape[0], dtype=np.uint16)

    #Build a cKdTree on the labels
    labelKdTree = cKDTree(labels)

    #Get the cloud point array
    cloudPointArray = np.asarray(pointCloud.getPoints(), dtype=np.float32)
    cloudDepthIndexArray = np.asarray(pointCloud.getIndices(), dtype=np.int32)

    cloudDepthYindices = cloudDepthIndexArray[:, 0].flatten()
    cloudDepthXindices = cloudDepthIndexArray[:, 1].flatten()

    #Find nearest neighbors of every point in the cloud point array
    _, labelInds = labelKdTree.query(cloudPointArray)
    labelInds = labelInds.astype(np.uint16)

    result[cloudDepthYindices, cloudDepthXindices] = labelInds

    return result

def handleSequence(argTuple):
    sequenceRoot, recordingName, destRoot = argTuple

    rootDirector = RootDirector(sequenceRoot)

    #Get the autolabelmanager for the sequence
    autoLabelManager = rootDirector.getAutoLabelFrameReadManager(recordingName)

    #Make destination directory if it doesn't exist
    destDir = os.path.join(destRoot, recordingName)
    os.makedirs(destDir)

    writer = DepthTrainingTFWriter(destDir)

    gc.collect()

    while True:
        #Extract the labeled frame, point cloud and labels
        label = autoLabelManager.getLabel()
        cloud = autoLabelManager.getCloud()

        #If the cloud has no points, skip it
        if (cloud.getPoints().size != 0):
            origDepthImg = autoLabelManager.getOrigFrame().getDepth()

            templateIndexImg = alternativeGetTemplateIndexImage(cloud, label)

            writer.add(origDepthImg, templateIndexImg)

        advanced = autoLabelManager.advance()
        if (not advanced):
            #No more frames annotated!
            break
    writer.flush()

if __name__ == '__main__':
    sequenceRoot = sys.argv[1]
    destRoot = sys.argv[2]
    rootDirector = RootDirector(sequenceRoot)
    recordingNames = rootDirector.getPeopleRecordingNames()

    #Great, now create a process for each source directory
    argsLists = []
    for recordingName in recordingNames:
        argsLists.append((sequenceRoot, recordingName, destRoot))
    p = Pool(NUM_PROCS)
    p.map(handleSequence, argsLists)


import numba

import autogalaxy as ag

from autolens.point.point_dataset import PointDataset
from autolens.point.point_solver import PointSolver
from autolens.point.fit_point.fluxes import FitFluxes
from autolens.point.fit_point.positions_image import FitPositionsImage
from autolens.point.fit_point.positions_source import FitPositionsSource
from autolens.lens.ray_tracing import Tracer

from autolens import exc

class FitPointDataset:
    def __init__(
        self, point_dataset: PointDataset, tracer: Tracer, point_solver: PointSolver
    ):

        self.point_dataset = point_dataset

        point_profile = tracer.extract_profile(profile_name=point_dataset.name)

        try:

            if isinstance(point_profile, ag.ps.PointSourceChi):

                self.positions = FitPositionsSource(
                    name=point_dataset.name,
                    positions=point_dataset.positions,
                    noise_map=point_dataset.positions_noise_map,
                    tracer=tracer,
                    point_profile=point_profile,
                )

            else:

                self.positions = FitPositionsImage(
                    name=point_dataset.name,
                    positions=point_dataset.positions,
                    noise_map=point_dataset.positions_noise_map,
                    point_solver=point_solver,
                    tracer=tracer,
                    point_profile=point_profile,
                )

        except exc.PointExtractionException:
            self.positions = None
        except (AttributeError, numba.errors.TypingError) as e:
            raise exc.FitException from e

        try:

            self.flux = FitFluxes(
                name=point_dataset.name,
                fluxes=point_dataset.fluxes,
                noise_map=point_dataset.fluxes_noise_map,
                positions=point_dataset.positions,
                tracer=tracer,
            )

        except exc.PointExtractionException:

            self.flux = None

    @property
    def log_likelihood(self) -> float:

        log_likelihood_positions = (
            self.positions.log_likelihood if self.positions is not None else 0.0
        )
        log_likelihood_flux = self.flux.log_likelihood if self.flux is not None else 0.0

        return log_likelihood_positions + log_likelihood_flux


import argparse
import copy
import math
import os
from itertools import chain
import random
import time
from collections import deque

import numpy as np
import tensorboardX
import torch
from torch import nn
import torch.nn.functional as F
import tqdm
import gin

from . import device, learning, evaluation, replay, augmentations
from . import learning_utils as lu

@gin.configurable
def super_sac(
    # required args
    agent,
    buffer,
    train_env,
    test_env,
    # compute kwargs
    bc_warmup_steps=0,
    num_steps_offline=1_000_000,
    num_steps_online=100_000,
    afbc_actor_updates_per_step=1,
    pg_actor_updates_per_step=1,
    critic_updates_per_step=1,
    target_critic_ensemble_n=2,
    batch_size=512,
    reuse_replay_dicts=True,
    # optimization kwargs
    actor_lr=3e-4,
    critic_lr=3e-4,
    encoder_lr=1e-4,
    alpha_lr=1e-4,
    actor_clip=40.0,
    critic_clip=40.0,
    encoder_clip=40.0,
    actor_l2=0.0,
    critic_l2=0.0,
    encoder_l2=0.0,
    pop=True,
    dr3_coeff=0.0,
    #   state abstraction kwargs
    markov_abstraction_updates_per_step=0,
    inverse_markov_coeff=0.0,
    contrastive_markov_coeff=0.0,
    smoothness_markov_coeff=0.0,
    smoothness_markov_max_dist=0.0,
    # rl kwargs
    use_exploration_process=False,
    exploration_param_init=1.0,
    exploration_param_final=0.1,
    exploration_param_anneal=1_000_000,
    exploration_update_clip=0.3,
    rolling_encoder_exploration=True,
    init_alpha=0.1,
    target_entropy_mul=1.0,
    gamma=0.99,
    mlp_tau=0.005,
    encoder_tau=0.01,
    target_delay=2,
    n_step=1,
    use_pg_update_online=True,
    use_afbc_update_online=True,
    weighted_bellman_temp=20.0,
    random_warmup_steps=0,
    weight_type="softmax",
    afbc_per=True,
    # env and eval kwargs
    transitions_per_online_step=1,
    infinite_bootstrap=True,
    ignore_all_dones=False,
    max_episode_steps=1_000_000,
    eval_episodes=10,
    eval_interval=10_000,
    render=False,
    # data augmentation
    augmenter=None,
    encoder_lambda=0.0,
    actor_lambda=0.0,
    aug_mix=0.9,
    # logging, misc
    logging_method="tensorboard",
    wandb_entity=os.getenv("SSAC_WANDB_ACCOUNT"),
    wandb_project=os.getenv("SSAC_WANDB_PROJECT"),
    base_save_path=os.getenv("SSAC_SAVE"),
    name="afbc_run",
    log_to_disk=True,
    log_interval=5000,
    hparams_config={},
    save_to_disk=True,
    save_best="eval/mean_return",
    evaluation_method=evaluation.evaluate_agent,
    verbosity=0,
):
    def _get_parallel_envs(env):
        _env = env
        while hasattr(_env, "env"):
            if hasattr(_env, "_PARALLEL_ACTORS"):
                return _env._PARALLEL_ACTORS
            else:
                _env = _env.env
        return 1

    num_envs = _get_parallel_envs(train_env)
    num_eval_envs = _get_parallel_envs(test_env)
    best_eval = -float("inf")

    if save_to_disk or log_to_disk:
        save_dir = make_process_dirs(name, base_save_path)
    if log_to_disk:
        hparams_config["ssac_run_name"] = name
        if logging_method == "tensorboard":
            writer = tensorboardX.SummaryWriter(save_dir)
            writer.add_hparams(hparams_config, {})
        elif logging_method == "wandb":
            assert (
                wandb_project is not None and wandb_entity is not None
            ), "Set `wandb_entity` and `wandb_project` kwargs. Note that you can also \n\
                set the os environment variables `SSAC_WANDB_ACCOUNT` and `SSAC_WANDB_PROJECT`, \n\
                respectively. Super SAC will default to those values."
            import wandb

            wandb.init(
                project=wandb_project, entity=wandb_entity, dir=save_dir, reinit=True
            )
            wandb.run.name = name
            wandb.config.update(hparams_config)
            wandb.run.save()

    if augmenter is None:
        augmenter = augmentations.AugmentationSequence(
            [augmentations.IdentityAug(batch_size)]
        )

    qprint = lambda x: print(x) if verbosity else None
    qprint(" ----- Super SAC -----")
    qprint(f"\tART: {agent.popart[0] is not False}")
    qprint(f"\tPOP: {pop}")
    qprint(
        f"\tBellman Backup Weight Type: {weight_type}; Temperature: {weighted_bellman_temp}"
    )
    qprint(f"\tBC Warmup Steps: {bc_warmup_steps}")
    qprint(f"\tEnsemble Size: {agent.ensemble_size}")
    qprint(f"\tCritic Ensemble Size: {agent.num_critics}")
    if target_critic_ensemble_n > agent.num_critics:
        target_critic_ensemble_n = agent.num_critics
        qprint("\t\tWarning: too many redq target agent critics. Overriding.")
    qprint(f"\tTD Target Critic Ensemble Size: {target_critic_ensemble_n}")
    qprint(f"\tCritic Updates per Step: {critic_updates_per_step}")
    qprint(f"\tDiscrete Actions: {agent.discrete}")
    qprint(f"\tActor Updates per Online Step: {pg_actor_updates_per_step}")
    qprint(f"\tActor Updates per Offline Step: {afbc_actor_updates_per_step}")
    qprint(f"\tQ-Value Uncertainty Exploration Bonus: {agent.ucb_bonus}")
    qprint(f"\tEncoder Lambda: {encoder_lambda}")
    qprint(f"\tActor Lambda: {actor_lambda}")
    qprint(f"\tUse PG Update Online: {use_pg_update_online}")
    qprint(f"\tUse BC Update Online: {use_afbc_update_online}")
    qprint(f"\tUse Random Exploration Noise: {use_exploration_process}")
    qprint(f"\tInit Alpha: {init_alpha}, Alpha LR: {alpha_lr}")
    qprint(f"\tDR3 Regularization Term: {dr3_coeff}")
    qprint(
        f"\tMarkov Abstraction Terms:\n \t\tInverse: {inverse_markov_coeff}, Contrastive: {contrastive_markov_coeff}, Smoothness: {smoothness_markov_coeff} with Max Dist {smoothness_markov_max_dist} "
    )
    qprint(f"\tAugmenter: {augmenter}")
    qprint(f"\tAug Mix: {aug_mix}")
    qprint(
        f"\tUsing Beta Dist: {not agent.discrete and agent.actors[0].dist_impl == 'beta'}"
    )
    qprint(f"\tParallel Training Envs: {num_envs}")
    qprint(f"\tParallel Eval Envs: {num_eval_envs}")
    qprint(" -----           -----")

    ###########
    ## SETUP ##
    ###########
    agent.to(device)
    agent.train()

    # optimizers
    critic_optimizer = torch.optim.Adam(
        chain(*(critic.parameters() for critic in agent.critics)),
        lr=critic_lr,
        weight_decay=critic_l2,
        betas=(0.9, 0.999),
    )
    offline_actor_optimizer = torch.optim.Adam(
        chain(*(actor.parameters() for actor in agent.actors)),
        lr=actor_lr,
        weight_decay=actor_l2,
        betas=(0.9, 0.999),
    )
    encoder_actorloss_optimizer = torch.optim.Adam(
        agent.encoder.parameters(),
        lr=encoder_lr,
        weight_decay=actor_l2,
        betas=(0.9, 0.999),
    )
    online_actor_optimizer = torch.optim.Adam(
        chain(*(actor.parameters() for actor in agent.actors)),
        lr=actor_lr,
        weight_decay=actor_l2,
        betas=(0.9, 0.999),
    )
    encoder_criticloss_optimizer = torch.optim.Adam(
        agent.encoder.parameters(),
        lr=encoder_lr,
        weight_decay=encoder_l2,
        betas=(0.9, 0.999),
    )
    encoder_markov_optimizer = torch.optim.Adam(
        chain(
            agent.encoder.parameters(),
            agent.inverse_model.parameters(),
            agent.contrastive_model.parameters(),
        ),
        lr=encoder_lr,
        weight_decay=encoder_l2,
        betas=(0.9, 0.999),
    )

    # max entropy, disabled with init_alpha = 0, alpha_lr = 0
    init_alpha = max(init_alpha, 1e-15)
    log_alphas = []
    log_alpha_optimizers = []
    for _ in range(agent.ensemble_size):
        log_alpha = torch.Tensor([math.log(init_alpha)]).to(device)
        log_alpha.requires_grad = True
        log_alphas.append(log_alpha)
        log_alpha_optimizers.append(
            torch.optim.Adam([log_alpha], lr=alpha_lr, betas=(0.5, 0.999))
        )
    if agent.discrete:
        target_entropy = -math.log(1.0 / train_env.action_space.n) * 0.98
    else:
        target_entropy = -train_env.action_space.shape[0]
    target_entropy *= target_entropy_mul

    # manual exploration, if applicable
    if use_exploration_process:
        if agent.discrete:
            random_process = lu.EpsilonGreedyExplorationNoise(
                action_space=train_env.action_space,
                eps_start=exploration_param_init,
                eps_final=exploration_param_final,
                steps_annealed=exploration_param_anneal,
            )
        else:
            random_process = lu.GaussianExplorationNoise(
                action_space=train_env.action_space,
                start_scale=exploration_param_init,
                final_scale=exploration_param_final,
                steps_annealed=exploration_param_anneal,
            )
    else:
        random_process = None

    ###################
    ## TRAINING LOOP ##
    ###################
    total_steps = bc_warmup_steps + num_steps_offline + num_steps_online
    progress_bar = lambda *x: tqdm.tqdm(range(*x)) if verbosity else range(*x)

    # warmup empty replay buffer
    if random_warmup_steps:
        lu.warmup_buffer(
            buffer=buffer,
            env=train_env,
            warmup_steps=random_warmup_steps,
            infinite_bootstrap=infinite_bootstrap,
            ignore_all_dones=ignore_all_dones,
            max_episode_steps=max_episode_steps,
            n_step=n_step,
            gamma=gamma,
            num_envs=num_envs,
        )

    done = True  # reset the env on first step
    exp_deque = deque([], maxlen=n_step)  # holds n-step transitions
    for step in progress_bar(total_steps):
        bc_logs, actor_logs, critic_logs, markov_logs = {}, {}, {}, {}

        ###############################
        ## Behavioral Cloning Update ##
        ###############################

        if step < bc_warmup_steps:
            if step == 0:
                qprint("[Behavioral Cloning]")
            bc_logs.update(
                learning.offline_actor_update(
                    buffer=buffer,
                    agent=agent,
                    actor_optimizer=offline_actor_optimizer,
                    actor_clip=actor_clip,
                    encoder_optimizer=encoder_actorloss_optimizer,
                    encoder_clip=encoder_clip,
                    update_encoder=True,
                    batch_size=batch_size,
                    augmenter=augmenter,
                    actor_lambda=actor_lambda,
                    aug_mix=aug_mix,
                    per=False,
                    discrete=agent.discrete,
                    filter_=False,
                )
            )
            bc_logs.update({"schedule/bc_actor_update": 1.0})
        else:
            bc_logs.update({"schedule/bc_actor_update": 0.0})

        # create target networks
        if step == bc_warmup_steps:
            qprint("[Creating Target Networks]")
            # create target networks (after behavioral cloning)
            target_agent = copy.deepcopy(agent)
            target_agent.to(device)
            for target_critic, agent_critic in zip(target_agent.critics, agent.critics):
                lu.hard_update(target_critic, agent_critic)
            lu.hard_update(target_agent.encoder, agent.encoder)
            target_agent.train()

        if step > bc_warmup_steps + num_steps_offline:
            #############################
            ## Environment Interaction ##
            #############################

            if step == bc_warmup_steps + num_steps_offline + 1:
                qprint("[Collecting Experience For the First Time]")
            for _ in range(transitions_per_online_step):
                if done:
                    state = train_env.reset()
                    steps_this_ep = 0
                    done = False
                    exp_deque.clear()
                    agent.encoder.reset_rolling()
                action = agent.sample_action(
                    state,
                    from_cpu=True,
                    num_envs=num_envs,
                    rolling=rolling_encoder_exploration,
                )
                if use_exploration_process:
                    actor_logs["exploration_noise_param"] = random_process.current_scale
                    action = random_process.sample(action, update_schedule=True)
                next_state, reward, done, info = train_env.step(action)
                if ignore_all_dones or (
                    infinite_bootstrap and steps_this_ep + 1 == max_episode_steps
                ):
                    # override the replay buffer version of done to False
                    buffer_done = (
                        np.expand_dims(np.array([False for _ in range(num_envs)]), 1)
                        if num_envs > 1
                        else False
                    )
                else:
                    buffer_done = done
                # put this transition in our n-step queue
                exp_deque.append((state, action, reward, next_state, buffer_done))
                if num_envs > 1:
                    done = done.any()
                if len(exp_deque) == exp_deque.maxlen:
                    # enough transitions to compute n-step returns
                    s, a, r, s1, d = exp_deque.popleft()
                    for i, trans in enumerate(exp_deque):
                        *_, r_i, s1, d = trans
                        r += (gamma ** (i + 1)) * r_i
                    # buffer gets n-step transition
                    buffer.push(s, a, r, s1, done=d, terminate_traj=done)
                state = next_state
                steps_this_ep += 1
                if steps_this_ep >= max_episode_steps:
                    done = True

        if step > bc_warmup_steps:
            ###################
            ## Critic Update ##
            ###################

            if step == bc_warmup_steps + 1:
                qprint("[First Critic Update]")
            critic_logs.update({"schedule/critic_update": 0.0})
            for critic_update in range(critic_updates_per_step):
                critic_logs, premade_replay_dicts = learning.critic_update(
                    buffer=buffer,
                    agent=agent,
                    target_agent=target_agent,
                    critic_optimizer=critic_optimizer,
                    encoder_optimizer=encoder_criticloss_optimizer,
                    log_alphas=log_alphas,
                    batch_size=batch_size,
                    gamma=gamma ** n_step,
                    critic_clip=critic_clip,
                    encoder_clip=encoder_clip,
                    target_critic_ensemble_n=target_critic_ensemble_n,
                    weighted_bellman_temp=weighted_bellman_temp,
                    weight_type=weight_type,
                    pop=pop,
                    augmenter=augmenter,
                    encoder_lambda=encoder_lambda,
                    aug_mix=aug_mix,
                    discrete=agent.discrete,
                    random_process=random_process,
                    noise_clip=exploration_update_clip,
                    per=False,
                    update_priorities=step < bc_warmup_steps + num_steps_offline
                    or use_afbc_update_online,
                    dr3_coeff=dr3_coeff,
                )
                critic_logs.update({"schedule/critic_update": 1.0})

                # move target model towards training model
                if (critic_update + step) % target_delay == 0:
                    for agent_critic, target_critic in zip(
                        agent.critics, target_agent.critics
                    ):
                        lu.soft_update(target_critic, agent_critic, mlp_tau)
                    lu.soft_update(target_agent.encoder, agent.encoder, encoder_tau)

        markov_logs.update({"schedule/markov_abstraction_update": 0.0})
        if step > bc_warmup_steps:
            for _ in range(markov_abstraction_updates_per_step):
                ##############################################
                ## Self-Supervised State Abstraction Update ##
                ##############################################

                markov_logs.update(
                    learning.markov_state_abstraction_update(
                        buffer=buffer,
                        agent=agent,
                        optimizer=encoder_markov_optimizer,
                        batch_size=batch_size,
                        augmenter=augmenter,
                        aug_mix=aug_mix,
                        discrete=agent.discrete,
                        inverse_coeff=inverse_markov_coeff,
                        contrastive_coeff=contrastive_markov_coeff,
                        smoothness_coeff=smoothness_markov_coeff,
                        smoothness_max_dist=smoothness_markov_max_dist,
                        grad_clip=encoder_clip,
                    )
                )
                markov_logs.update({"schedule/markov_abstraction_update": 1.0})

        if (step > bc_warmup_steps and step < bc_warmup_steps + num_steps_offline) or (
            step >= bc_warmup_steps + 1 + num_steps_offline and use_afbc_update_online
        ):
            #######################
            ## AWAC Actor Update ##
            #######################

            if step == bc_warmup_steps + 1:
                qprint("[First Offline Actor Update]")
            if step == bc_warmup_steps + num_steps_offline:
                qprint("[First Online Filtered BC Update]")
            for actor_update in range(afbc_actor_updates_per_step):
                _use_past_dicts = (
                    not afbc_per and actor_update == 0
                ) and reuse_replay_dicts
                actor_logs.update(
                    learning.offline_actor_update(
                        buffer=buffer,
                        agent=agent,
                        actor_optimizer=offline_actor_optimizer,
                        # train encoder with critic grads only
                        encoder_optimizer=encoder_actorloss_optimizer,
                        encoder_clip=encoder_clip,
                        update_encoder=False,
                        batch_size=batch_size,
                        actor_clip=actor_clip,
                        augmenter=augmenter,
                        actor_lambda=actor_lambda,
                        aug_mix=aug_mix,
                        per=afbc_per,
                        premade_replay_dicts=premade_replay_dicts
                        if _use_past_dicts
                        else None,
                        discrete=agent.discrete,
                        filter_=True,
                    )
                )
            actor_logs.update({"schedule/afbc_actor_update": 1.0})
        else:
            actor_logs.update({"schedule/afbc_actor_update": 0.0})

        if step > bc_warmup_steps + num_steps_offline and use_pg_update_online:
            ######################
            ## DPG Actor Update ##
            ######################

            if step == bc_warmup_steps + num_steps_offline + 1:
                qprint("[First Online Actor Update]")
            for actor_update in range(pg_actor_updates_per_step):
                _use_past_dicts = (actor_update == 0) and reuse_replay_dicts
                actor_logs.update(
                    learning.online_actor_update(
                        buffer=buffer,
                        agent=agent,
                        pop=pop,
                        actor_optimizer=online_actor_optimizer,
                        log_alphas=log_alphas,
                        batch_size=batch_size,
                        aug_mix=aug_mix,
                        clip=actor_clip,
                        augmenter=augmenter,
                        per=False,
                        discrete=agent.discrete,
                        random_process=random_process,
                        noise_clip=exploration_update_clip,
                        premade_replay_dicts=premade_replay_dicts
                        if _use_past_dicts
                        else None,
                        use_baseline=False,
                    )
                )
            actor_logs.update({"schedule/pg_actor_update": 1.0})
        else:
            actor_logs.update({"schedule/pg_actor_update": 0.0})

        if (
            step > bc_warmup_steps + num_steps_offline
            and init_alpha > 0
            and alpha_lr > 0
        ):
            ##################
            ## Alpha Update ##
            ##################

            if step == bc_warmup_steps + num_steps_offline + 1:
                qprint("[First Alpha Update]")
            _use_past_dicts = reuse_replay_dicts
            actor_logs.update(
                learning.alpha_update(
                    buffer=buffer,
                    agent=agent,
                    optimizers=log_alpha_optimizers,
                    batch_size=batch_size,
                    log_alphas=log_alphas,
                    augmenter=augmenter,
                    aug_mix=aug_mix,
                    target_entropy=target_entropy,
                    premade_replay_dicts=premade_replay_dicts
                    if _use_past_dicts
                    else None,
                    discrete=agent.discrete,
                )
            )
            actor_logs.update({"schedule/alpha_update": 1.0})
        else:
            actor_logs.update({"schedule/alpha_update": 0.0})

        #############
        ## LOGGING ##
        #############
        if (step % log_interval == 0) and log_to_disk:
            performance_logs = {
                "replay_buffer_total_samples": buffer.total_sample_calls,
                "len(buffer)": len(buffer),
            }
            if logging_method == "tensorboard":
                for key, val in critic_logs.items():
                    writer.add_scalar(key, val, step)
                for key, val in actor_logs.items():
                    writer.add_scalar(key, val, step)
                for key, val in bc_logs.items():
                    writer.add_scalar(key, val, step)
                for key, val in performance_logs.items():
                    writer.add_scalar(key, val, step)
                for key, val in markov_logs.items():
                    writer.add_scalar(key, val, step)
            elif logging_method == "wandb":
                wandb.log(critic_logs, step=step)
                wandb.log(actor_logs, step=step)
                wandb.log(bc_logs, step=step)
                wandb.log(performance_logs, step=step)
                wandb.log(markov_logs, step=step)

        if (
            (step % eval_interval == 0) or (step == total_steps - 1)
        ) and eval_interval > 0:
            eval_logs = evaluation_method(
                agent,
                test_env,
                eval_episodes,
                max_episode_steps,
                render,
                verbosity=verbosity,
                num_envs=num_eval_envs,
            )
            if log_to_disk:
                accepted_exp_pct = lu.compute_filter_stats(
                    buffer=buffer,
                    agent=agent,
                    augmenter=augmenter,
                    batch_size=batch_size,
                )
                if logging_method == "tensorboard":
                    for key, val in eval_logs.items():
                        writer.add_scalar(key, val, step)
                    writer.add_scalar("Accepted Exp Pct", accepted_exp_pct, step)
                elif logging_method == "wandb":
                    wandb.log({**eval_logs, **{"Accepted Exp Pct": accepted_exp_pct}})

            if save_to_disk:
                if save_best is not None:
                    assert (
                        save_best in eval_logs
                    ), f"`save_best` metric {save_best} not found in evaluation logs with keys {eval_logs.keys()}"
                    eval_metric = eval_logs[save_best]
                    if eval_metric > best_eval:
                        qprint(f"[Saving Agent with {save_best}: {eval_metric}]")
                        best_eval = eval_metric
                        agent.save(save_dir)
                else:
                    agent.save(save_dir)

    if save_to_disk:
        agent.save(save_dir)

    return agent

def make_process_dirs(run_name, base_path):
    if base_path is None:
        base_path = "./saves"
    base_dir = os.path.join(base_path, run_name)
    i = 0
    while os.path.exists(base_dir + f"_{i}"):
        i += 1
    base_dir += f"_{i}"
    os.makedirs(base_dir)
    return base_dir


import argparse
import math
import random
import os
import copy

os.environ["CUDA_VISIBLE_DEVICES"]="0"

from numpy.core.fromnumeric import resize
import dnnlib

import numpy as np
import torch
from torch import nn, autograd, optim
from torch.nn import functional as F
from torch.utils import data
import torch.distributed as dist
from torchvision import transforms, utils
from tqdm import tqdm
from torch_utils import image_transforms
from PIL import Image

from model import Generator, G_NET, Encoder, UNet
from finegan_config import finegan_config

# try:
import wandb

# except ImportError:
#     wandb = None

from op import conv2d_gradfix
from dataset import MultiResolutionDataset

from distributed import (
    get_rank,
    synchronize,
    reduce_loss_dict,
    reduce_sum,
    get_world_size,
)
# from non_leaking import augment, AdaptiveAugment

# os.environ["CUDA_VISIBLE_DEVICES"]="0"

# os.environ['PYTHONWARNINGS'] = 'ignore:semaphore_tracker:UserWarning'

def weights_init(m):
    classname = m.__class__.__name__
    if isinstance(m, nn.Conv2d):
        nn.init.orthogonal_(m.weight.data, 1.0)
    elif isinstance(m, nn.BatchNorm2d):
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)
    elif isinstance(m, nn.Linear):
        nn.init.orthogonal_(m.weight.data, 1.0)
        if m.bias is not None:
            m.bias.data.fill_(0.0)
    elif classname == 'MnetConv':
        nn.init.constant_(m.mask_conv.weight.data, 1)

def data_sampler(dataset, shuffle, distributed):
    if distributed:
        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)

    if shuffle:
        return data.RandomSampler(dataset)

    else:
        return data.SequentialSampler(dataset)

def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag

def accumulate(model1, model2, decay=0.999):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())

    for k in par1.keys():
        par1[k].data.mul_(decay).add_(par2[k].data, alpha=1 - decay)

def sample_data(loader):
    while True:
        for batch in loader:
            yield batch

def d_logistic_loss(real_pred, fake_pred):
    real_loss = F.softplus(-real_pred)
    fake_loss = F.softplus(fake_pred)

    return real_loss.mean() + fake_loss.mean()

def d_r1_loss(real_pred, real_img):
    with conv2d_gradfix.no_weight_gradients():
        grad_real, = autograd.grad(
            outputs=real_pred.sum(), inputs=real_img, create_graph=True
        )
    grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean()

    return grad_penalty

def g_nonsaturating_loss(fake_pred):
    loss = F.softplus(-fake_pred).mean()
    return loss

def g_path_regularize(fake_img, latents, mean_path_length, decay=0.01):
    noise = torch.randn_like(fake_img) / math.sqrt(
        fake_img.shape[2] * fake_img.shape[3]
    )
    grad, = autograd.grad(
        outputs=(fake_img * noise).sum(), inputs=latents, create_graph=True
    )

    path_lengths = torch.sqrt(grad.pow(2).sum(2).mean(1))

    path_mean = mean_path_length + decay * (path_lengths.mean() - mean_path_length)

    path_penalty = (path_lengths - path_mean).pow(2).mean()

    return path_penalty, path_mean.detach(), path_lengths

def set_grad_none(model, targets):
    for n, p in model.named_parameters():
        if n in targets:
            p.grad = None

def make_noise(batch, latent_dim, n_noise, device):
    if n_noise == 1:
        return torch.randn(batch, latent_dim, device=device)
    noises = torch.randn(n_noise, batch, latent_dim, device=device)
    return noises

def mixing_noise(batch, latent_dim, prob, device):
    prob = 1
    if prob > 0 and random.random() < prob:
        return make_noise(batch, latent_dim, 2, device)
    else:
        return [make_noise(batch, latent_dim, 1, device)]

def child_to_parent(c_code, c_dim, p_dim):
    ratio = c_dim / p_dim
    cid = torch.argmax(c_code,  dim=1)
    pid = (cid / ratio).long()
    p_code = torch.zeros([c_code.size(0), p_dim], device=c_code.device)
    for i in range(c_code.size(0)):
        p_code[i][pid[i]] = 1
    return p_code

def sample_codes(batch, z_dim, b_dim, p_dim, c_dim, device):
    z = torch.randn(batch, z_dim, device=device)
    c = torch.zeros(batch, c_dim, device=device)
    cid = np.random.randint(c_dim, size=batch)
    for i in range(batch):
        c[i, cid[i]] = 1

    p = child_to_parent(c, c_dim, p_dim)
    b = c.clone()
    return z, b, p, c

def rand_sample_codes(prev_z=None, prev_b=None, prev_p=None, prev_c=None, rand_code=['b', 'p']):
    '''
    rand code default: keeping z and c
    '''

    if prev_z is not None:
        device = prev_z.device
        batch = prev_z.size(0)
    elif prev_b is not None:
        device = prev_b.device
        batch = prev_b.size(0)
    elif prev_p is not None:
        device = prev_p.device
        batch = prev_p.size(0)
    elif prev_c is not None:
        device = prev_c.device
        batch = prev_c.size(0)
    else:
        sys.exit(0)

    if 'z' in rand_code:
        z = torch.randn(batch, prev_z.size(1), device=device)
    else:
        z = prev_z

    if 'b' in rand_code:
        b = torch.zeros(batch, prev_b.size(1), device=device)
        bid = np.random.randint(prev_b.size(1), size=batch)
        for i in range(batch):
            b[i, bid[i]] = 1
    else:
        b = prev_b

    if 'p' in rand_code:
        p = torch.zeros(batch, prev_p.size(1), device=device)
        pid = np.random.randint(prev_p.size(1), size=batch)
        for i in range(batch):
            p[i, pid[i]] = 1
    else:
        p = prev_p

    if 'c' in rand_code:
        c = torch.zeros(batch, prev_c.size(1), device=device)
        cid = np.random.randint(prev_c.size(1), size=batch)
        for i in range(batch):
            c[i, cid[i]] = 1
    else:
        c = prev_c

    return z, b, p, c

def binarization_loss(mask):
    return torch.min(1-mask, mask).mean()

def train(args, fine_generator, style_generator, mpnet, mknet, mp_optim, mk_optim, device):

    pbar = range(args.iter)

    if get_rank() == 0:
        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0.01)

    # cur_nimg = 0
    loss_dict = {}

    if args.distributed:
        mp_module = mpnet.module
        mk_module = mknet.module
        fine_g_module = fine_generator.module
        style_g_module = style_generator.module
    else:
        mp_module = mpnet
        mk_module = mknet
        fine_g_module = fine_generator
        style_g_module = style_generator

    # criterion_construct = nn.MSELoss()
    # criterion_reconstruct = nn.MSELoss().to(device)

    if args.trunc:
        truncation = 0.7
        trunc = style_generator.mean_latent(4096).detach()
        trunc.requires_grad_(False)
    else:
        truncation = 1
        trunc = None

    style_generator.eval()
    style_generator.requires_grad_(False)
    fine_generator.eval()
    fine_generator.requires_grad_(False)

    for idx in pbar:
        i = idx + args.start_iter

        if i > args.iter:
            print("Done!")
            break

        mpnet.train()
        mknet.train()

        # ############# train mk network #############
        # mknet.requires_grad_(True)
        # # mpnet.requires_grad_(False)

        # z, b, p, c = sample_codes(args.batch, args.z_dim, args.b_dim, args.p_dim, args.c_dim, device)
        # if not args.tie_code:
        #     z, b, p, c = rand_sample_codes(prev_z=z, prev_b=b, prev_p=p, prev_c=c, rand_code=['b', 'p'])

        # fine_img, mask = fine_generator(z, b, p, c, rtn_mk=True)
        # pred_mask = mknet(fine_img)

        # bin_loss = binarization_loss(pred_mask) * args.bin
        # mk_loss = F.mse_loss(pred_mask, mask) * args.mk

        # mknet_loss = mk_loss + bin_loss
        loss_dict["mk"] = torch.zeros(1).to(device)
        # loss_dict["bin"] = bin_loss / args.bin

        # mknet.zero_grad()
        # mknet_loss.backward()
        # mk_optim.step()

        ############# train mapping network #############
        mpnet.requires_grad_(True)

        noise = mixing_noise(args.batch, args.latent, args.mixing, device)

        style_img, latent = style_generator(noise, return_latents=True, randomize_noise=False)
        _style_img = F.interpolate(style_img, size=(128, 128), mode='area')

        #########
        with torch.no_grad():
            mask = mknet(_style_img)

        grad, = autograd.grad(
            outputs=(_style_img * noise).sum(), inputs=latent, create_graph=True
        )

        ##########

        wp_code = mpnet(_style_img)

        mp_loss = F.mse_loss(wp_code, torch.transpose(noise, 0, 1)) * args.mp
        loss_dict["mp"] = mp_loss / args.mp

        mpnet.zero_grad()
        mp_loss.backward()
        mp_optim.step()

        ############# ############# #############
        loss_reduced = reduce_loss_dict(loss_dict)
        mp_loss_val = loss_reduced["mp"].mean().item()
        mk_loss_val = loss_reduced["mk"].mean().item()
        # bin_loss_val = loss_reduced["bin"].mean().item()

        if get_rank() == 0:
            pbar.set_description(
                (
                    f"mp: {mp_loss_val:.4f}; mk: {mk_loss_val:.4f}"
                )
            )

            if wandb and args.wandb:
                wandb.log(
                    {
                        "MP": mp_loss_val,
                        "MK": mk_loss_val,
                    }
                )

            # if i % 500 == 0:
            #     with torch.no_grad():
            #         mpnet.eval()
            #         mknet.eval()

            #         z, b, p, c = sample_codes(args.n_sample, args.z_dim, args.b_dim, args.p_dim, args.c_dim, device)
            #         if not args.tie_code:
            #             z, b, p, c = rand_sample_codes(prev_z=z, prev_b=b, prev_p=p, prev_c=c, rand_code=['b', 'p'])

            #         fine_img, _ = fine_generator(z, b, p, c)
            #         wp_code = mpnet(fine_img)
            #         rec_fine, _ = style_generator([wp_code], input_is_latent=True)

            #         _rec_fine = F.interpolate(rec_fine, size=(128, 128), mode='area')

            #         rec_mk = mknet(_rec_fine)

            #         noise = mixing_noise(args.n_sample, args.latent, args.mixing, device)
            #         style_img, _ = style_generator(noise, return_latents=True)
            #         _style_img = F.interpolate(style_img, size=(128, 128), mode='area')
            #         wp_code = mpnet(_style_img)
            #         rec_style, _ = style_generator([wp_code], input_is_latent=True)

            #         utils.save_image(
            #             fine_img,
            #             f"sample/{str(i).zfill(6)}_0.png",
            #             nrow=8,
            #             normalize=True,
            #             range=(-1, 1),
            #         )

            #         utils.save_image(
            #             rec_fine,
            #             f"sample/{str(i).zfill(6)}_1.png",
            #             nrow=8,
            #             normalize=True,
            #             range=(-1, 1),
            #         )

            #         utils.save_image(
            #             rec_mk,
            #             f"sample/{str(i).zfill(6)}_2.png",
            #             nrow=8,
            #             normalize=True,
            #             range=(0, 1),
            #         )

            #         utils.save_image(
            #             style_img,
            #             f"sample/{str(i).zfill(6)}_3.png",
            #             nrow=8,
            #             normalize=True,
            #             range=(-1, 1),
            #         )

            #         utils.save_image(
            #             rec_style,
            #             f"sample/{str(i).zfill(6)}_4.png",
            #             nrow=8,
            #             normalize=True,
            #             range=(-1, 1),
            #         )
            #         if wandb and args.wandb:
            #             wandb.log(
            #                 {
            #                     "fine image": [wandb.Image(Image.open(f"sample/{str(i).zfill(6)}_0.png").convert("RGB"))],
            #                     "recon fine": [wandb.Image(Image.open(f"sample/{str(i).zfill(6)}_1.png").convert("RGB"))],
            #                     "recon mask": [wandb.Image(Image.open(f"sample/{str(i).zfill(6)}_2.png").convert("RGB"))],
            #                     "style image": [wandb.Image(Image.open(f"sample/{str(i).zfill(6)}_3.png").convert("RGB"))],
            #                     "recon style": [wandb.Image(Image.open(f"sample/{str(i).zfill(6)}_4.png").convert("RGB"))],
            #                 }
            #             )

            if i % 40000 == 0 and i != args.start_iter:
                torch.save(
                    {
                        "style_g": style_g_module.state_dict(),
                        "fine": fine_g_module.state_dict(),
                        "mp": mp_module.state_dict(),
                        "mk": mk_module.state_dict(),
                        "mp_optim": mp_optim.state_dict(),
                        "mk_optim": mk_optim.state_dict(),
                        "args": args,
                        "cur_iter": i,
                    },
                    f"checkpoint/{str(i).zfill(6)}_1_.pt",
                )

if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser(description="mpnet trainer")

    parser.add_argument('--arch', type=str, default='stylegan2', help='model architectures (stylegan2 | swagan)')
    parser.add_argument(
        "--iter", type=int, default=800000, help="total training iterations"
    )
    parser.add_argument(
        "--batch", type=int, default=16, help="batch sizes for each gpus"
    )
    parser.add_argument(
        "--n_sample",
        type=int,
        default=8,
        help="number of the samples generated during training",
    )
    parser.add_argument(
        "--size", type=int, default=256, help="image sizes for the model"
    )
    parser.add_argument(
        "--mixing", type=float, default=0.9, help="probability of latent code mixing"
    )
    parser.add_argument(
        "--ckpt",
        type=str,
        default=None,
        help="path to previous trained checkpoint",
    )
    parser.add_argument(
        "--style_model",
        type=str,
        default=None,
        help="path to stylegan",
    )
    parser.add_argument(
        "--fine_model",
        type=str,
        default=None,
        help="path to finegan",
    )
    parser.add_argument("--lr_mp", type=float, default=2e-3, help="mapping network learning rate")
    # parser.add_argument("--lr_d", type=float, default=2e-5, help="discriminator learning rate")
    parser.add_argument(
        "--channel_multiplier",
        type=int,
        default=2,
        help="channel multiplier factor for the model. config-f = 2, else = 1",
    )
    parser.add_argument(
        "--wandb", action="store_true", help="use weights and biases logging"
    )
    parser.add_argument(
        "--local_rank", type=int, default=0, help="local rank for distributed training"
    )
    parser.add_argument("--d_reg_every", type=int, default=16)
    parser.add_argument("--r1", type=float, default=10)
    parser.add_argument(
        "--trunc", action="store_true", help="use truncation"
    )
    parser.add_argument('--mp_arch', type=str, default='encoder',
                        help='model architectures (vanilla | encoder)')
    parser.add_argument(
        "--tie_code", action="store_true", help="use tied codes"
    )
    parser.add_argument('--ds_name', type=str, default='STANFORDCAR',
                        help='dataset used for training finegan (LSUNCAR | CUB | STANFORDCAR)')

    ## weights
    # parser.add_argument("--adv", type=float, default=1, help="weight of the adv loss")
    # parser.add_argument("--mse", type=float, default=1e2, help="weight of the mse loss")
    parser.add_argument("--mp", type=float, default=1, help="weight of the latent recon loss")
    parser.add_argument("--mk", type=float, default=1e1, help="weight of mask recon")
    parser.add_argument("--bin", type=float, default=1, help="weight of mask recon")

    args = parser.parse_args()

    n_gpu = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
    args.distributed = n_gpu > 1

    if args.distributed:
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
        synchronize()

    args.latent = 512
    args.n_mlp = 8

    args.start_iter = 0

    args.z_dim = finegan_config[args.ds_name]['Z_DIM']
    args.b_dim = finegan_config[args.ds_name]['FINE_GRAINED_CATEGORIES']
    args.p_dim = finegan_config[args.ds_name]['SUPER_CATEGORIES']
    args.c_dim = finegan_config[args.ds_name]['FINE_GRAINED_CATEGORIES']

    if args.arch == 'stylegan2':
        from model import Generator, Discriminator

    elif args.arch == 'swagan':
        from swagan import Generator, Discriminator

    style_generator = Generator(
        size=args.size,
        style_dim=args.latent,
        n_mlp=args.n_mlp,
        channel_multiplier=args.channel_multiplier
    ).to(device)

    fine_generator = G_NET(ds_name=args.ds_name).to(device)

    if args.mp_arch == 'vanilla':
        mpnet = MappingNetwork(
            num_ws=style_generator.n_latent,
            w_dim=args.latent
        ).to(device)
    # https://github.com/bryandlee/stylegan2-encoder-pytorch
    elif args.mp_arch == 'encoder':
        mpnet = Encoder(
            size=128,
            num_ws=2,
            w_dim=args.latent
        ).to(device)

    mknet = UNet(
        n_channels = 3,
        n_classes = 1,
        bilinear = True,
    ).to(device)

    mp_optim = optim.Adam(
        mpnet.parameters(),
        lr=args.lr_mp,
        betas=(0, 0.99),
    )

    mk_optim = optim.Adam(
        mknet.parameters(),
        lr=args.lr_mp,
        betas=(0, 0.99),
    )

    if args.ckpt is not None:
        print("load model:", args.ckpt)

        ckpt = torch.load(args.ckpt, map_location=lambda storage, loc: storage)
        # train_args = ckpt['args']
        # args.start_iter = ckpt['cur_iter']

        # mpnet.load_state_dict(ckpt["mp"])
        # mknet.load_state_dict(ckpt["mk"])

        if args.style_model is None:
            style_generator.load_state_dict(ckpt["style_g"])
        #     mp_optim.load_state_dict(ckpt["mp_optim"])

        # if args.fine_model is None:
        #     fine_generator.load_state_dict(ckpt["fine"])
        #     mk_optim.load_state_dict(ckpt["mk_optim"])

    # if specify stylegan checkpoint, overwrite stylegan from ckpt
    if args.style_model is not None:
        print("load style model:", args.style_model)
        style_dict = torch.load(args.style_model, map_location=lambda storage, loc: storage)
        style_generator.load_state_dict(style_dict["g_ema"])
        # d_optim.load_state_dict(style_dict["d_optim"])

    # if specify finegan checkpoint, overwrite finegan from ckpt
    if args.fine_model is not None:
        print("load fine model:", args.fine_model)
        fine_dict = torch.load(args.fine_model, map_location=lambda storage, loc: storage)
        fine_generator.load_state_dict(fine_dict)

    if args.distributed:
        style_generator = nn.parallel.DistributedDataParallel(
            style_generator,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
        )

        fine_generator = nn.parallel.DistributedDataParallel(
            fine_generator,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
        )

        mpnet = nn.parallel.DistributedDataParallel(
            mpnet,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
        )

        mknet = nn.parallel.DistributedDataParallel(
            mknet,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
        )

    if get_rank() == 0 and wandb is not None and args.wandb:
        wandb.init(project="map net style distribute")

    train(args, fine_generator, style_generator, mpnet, mknet, mp_optim, mk_optim, device)


#!/usr/bin/python
# -*- coding: utf-8 -*-
# Author: violinsolo
# Created on 17/09/2018

from data.data_reader import DataReader

dr = DataReader(name='./../data/gtd')

df = dr.df
all_cols = ['eventid', 'iyear', 'imonth', 'iday', 'approxdate', 'extended', 'resolution', 'country', 'country_txt',
            'region', 'region_txt', 'provstate', 'city', 'latitude', 'longitude', 'specificity', 'vicinity', 'location',
            'summary', 'crit1', 'crit2', 'crit3', 'doubtterr', 'alternative', 'alternative_txt', 'multiple', 'success',
            'suicide', 'attacktype1', 'attacktype1_txt', 'attacktype2', 'attacktype2_txt', 'attacktype3',
            'attacktype3_txt', 'targtype1', 'targtype1_txt', 'targsubtype1', 'targsubtype1_txt', 'corp1', 'target1',
            'natlty1', 'natlty1_txt', 'targtype2', 'targtype2_txt', 'targsubtype2', 'targsubtype2_txt', 'corp2',
            'target2', 'natlty2', 'natlty2_txt', 'targtype3', 'targtype3_txt', 'targsubtype3', 'targsubtype3_txt',
            'corp3', 'target3', 'natlty3', 'natlty3_txt', 'gname', 'gsubname', 'gname2', 'gsubname2', 'gname3',
            'gsubname3', 'motive', 'guncertain1', 'guncertain2', 'guncertain3', 'individual', 'nperps', 'nperpcap',
            'claimed', 'claimmode', 'claimmode_txt', 'claim2', 'claimmode2', 'claimmode2_txt', 'claim3', 'claimmode3',
            'claimmode3_txt', 'compclaim', 'weaptype1', 'weaptype1_txt', 'weapsubtype1', 'weapsubtype1_txt',
            'weaptype2', 'weaptype2_txt', 'weapsubtype2', 'weapsubtype2_txt', 'weaptype3', 'weaptype3_txt',
            'weapsubtype3', 'weapsubtype3_txt', 'weaptype4', 'weaptype4_txt', 'weapsubtype4', 'weapsubtype4_txt',
            'weapdetail', 'nkill', 'nkillus', 'nkillter', 'nwound', 'nwoundus', 'nwoundte', 'property', 'propextent',
            'propextent_txt', 'propvalue', 'propcomment', 'ishostkid', 'nhostkid', 'nhostkidus', 'nhours', 'ndays',
            'divert', 'kidhijcountry', 'ransom', 'ransomamt', 'ransomamtus', 'ransompaid', 'ransompaidus', 'ransomnote',
            'hostkidoutcome', 'hostkidoutcome_txt', 'nreleased', 'addnotes', 'scite1', 'scite2', 'scite3', 'dbsource',
            'INT_LOG', 'INT_IDEO', 'INT_MISC', 'INT_ANY', 'related']

# get all data here.

target_col_name = [
    'eventid', 'iyear', 'imonth', 'iday', 'extended',
    'crit1', 'crit2', 'crit3', 'doubtterr', 'alternative', 'multiple',
    'country', 'region', 'provstate', 'city', 'specificity', 'vicinity',
    'attacktype1', 'attacktype2', 'attacktype3', 'success', 'suicide',
    'weaptype1', 'weapsubtype1', 'weaptype2', 'weapsubtype2', 'weaptype3', 'weapsubtype3', 'weaptype4', 'weapsubtype4',
    'targtype1', 'targsubtype1', 'targtype2', 'targsubtype2', 'targtype3', 'targsubtype3',
    'gname', 'nperps', 'nperpcap', 'claimed', 'claimmode', 'compclaim',
    'nkill', 'nkillus', 'nkillter', 'nwound', 'nwoundus', 'nwoundte', 'property', 'propextent', 'propvalue', 'ishostkid', 'nhostkid', 'nhostkidus', 'nhours', 'ndays', 'ransom', 'hostkidoutcome', 'nreleased',
    'INT_LOG', 'INT_IDEO', 'INT_MISC', 'INT_ANY', 'related'
]
df_target = df[target_col_name]

# fill nas
df_target.fillna({
    'eventid': 0, 'iyear': 0, 'imonth': 0, 'iday': 0, 'extended': 0,
    'crit1': 0, 'crit2': 0, 'crit3': 0, 'doubtterr': 0., 'alternative': 0., 'multiple': 0.,
    'country': 0, 'region': 0, 'provstate': '', 'city': '', 'specificity': 0., 'vicinity': 0,
    'attacktype1': 9, 'attacktype2': 9, 'attacktype3': 9, 'success': 0, 'suicide': 0,
    'weaptype1': 13, 'weapsubtype1': 27, 'weaptype2': 13, 'weapsubtype2': 27, 'weaptype3': 13, 'weapsubtype3': 27, 'weaptype4': 13, 'weapsubtype4': 27,
    'targtype1': 20, 'targsubtype1': 0, 'targtype2': 20, 'targsubtype2': 0, 'targtype3': 20, 'targsubtype3': 0,
    'gname': 'Unknown', 'nperps': 0, 'nperpcap': 0, 'claimed': 0, 'claimmode': 10, 'compclaim': 0,
    'nkill': 0, 'nkillus': 0, 'nkillter': 0, 'nwound': 0, 'nwoundus': 0, 'nwoundte': 0, 'property': -9, 'propextent': 4, 'propvalue': -99, 'ishostkid': 0, 'nhostkid': 0, 'nhostkidus': 0, 'nhours': 0, 'ndays': -99, 'ransom': 0, 'hostkidoutcome': 0, 'nreleased': 0,
    'INT_LOG': 0, 'INT_IDEO': 0, 'INT_MISC': 0, 'INT_ANY': 0, 'related': ''

}, inplace=True)

# do pre-process
df_needed_ori = df_target.drop(['eventid', 'iyear', 'imonth', 'iday', 'related'], axis=1)
df_1516_ori = df_target.loc[(df_target['iyear'] == 2015) | (df_target['iyear'] == 2016), :].drop(['eventid', 'iyear', 'imonth', 'iday', 'related'], axis=1)
df_17_ori = df_target.loc[df_target['iyear'] == 2017, :].drop(['eventid', 'iyear', 'imonth', 'iday', 'related'], axis=1)
df_needed = df_needed_ori[df_needed_ori['gname'] != 'Unknown']
df_needed_unk = df_needed_ori[df_needed_ori['gname'] == 'Unknown']
df_dbscan_train_unk = df_1516_ori[(df_1516_ori['gname'] == 'Unknown') & (df_1516_ori['claimed'] == 0)]
df_dbscan_train_real = df_dbscan_train_unk.drop(['gname'], axis=1)

# do labelencoding
from sklearn import preprocessing

# encode gname
le_gname = preprocessing.LabelEncoder()
en_gname = le_gname.fit_transform(df_needed['gname'])
df_needed['gname'] = en_gname.reshape((-1, 1))

# encode provstate
le_provstate = preprocessing.LabelEncoder()
en_provstate = le_provstate.fit_transform(df_needed['provstate'])
df_needed['provstate'] = en_provstate.reshape((-1, 1))
le_ds_provstate = preprocessing.LabelEncoder()
en_ds_provstate = le_ds_provstate.fit_transform(df_dbscan_train_real['provstate'])
df_dbscan_train_real['provstate'] = en_ds_provstate.reshape((-1, 1))

# encode city
le_city = preprocessing.LabelEncoder()
en_city = le_city.fit_transform(df_needed['city'])
df_needed['city'] = en_city.reshape((-1, 1))
le_ds_city = preprocessing.LabelEncoder()
en_ds_city = le_ds_city.fit_transform(df_dbscan_train_real['city'])
df_dbscan_train_real['city'] = en_ds_city.reshape((-1, 1))

# Filter all unknown gnames...
import pandas as pd
import numpy as np

allY = df_needed['gname']
allX = df_needed.drop(['gname'], axis=1)

from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier

print("====> Begin split...")
SEED = 4
x_train, x_test, y_train, y_test = train_test_split(allX.values, allY.values, test_size=0.25, random_state=SEED)
print("====> Split finished...")

feat_labels = allX.columns[0:]
forest = RandomForestClassifier(n_estimators=10, random_state=SEED, n_jobs=-1, verbose=2)
print("====> Training begin...")
forest.fit(x_train, y_train)
print("====> Training Finished...")

importances = forest.feature_importances_
indices = np.argsort(importances)[::-1]
for f in range(x_train.shape[1]):
    print("%2d) %-*s %f" % (f + 1, 30, feat_labels[indices[f]], importances[indices[f]]))

threshold = 0.005
x_selected = x_train[:, importances > threshold]
print(x_selected.shape)

all_tar_idxs = []
for f in range(x_train.shape[1]):
    if importances[indices[f]] > threshold:
        all_tar_idxs.append(feat_labels[indices[f]])

print(all_tar_idxs)

x_selected_fea = df_dbscan_train_real[all_tar_idxs]
print(x_selected_fea.shape)

# NOW Doing DBSCAN
import hdbscan
from sklearn.cluster import DBSCAN
import time

#DBSCAN
t0 = time.time()
# dbscan = DBSCAN(eps=100, min_samples=3)
clusterer = hdbscan.HDBSCAN(min_cluster_size=5, min_samples=3)
cluster_labels = clusterer.fit_predict(x_selected_fea)
# hierarchy = clusterer.cluster_hierarchy_
# alt_labels = hierarchy.get_clusters(0.100, 5)
# hierarchy.plot()

cunique = np.unique(cluster_labels)
l = {i: cluster_labels[cluster_labels==i].size for i in cunique}
l2 = [cluster_labels[cluster_labels==i].size for i in cunique]

# ordered dict for unique labels and it's count
dict = sorted(l.items(), key=lambda d:d[1], reverse = True)

# to store final cluster result..
df_hdbscan_final_sub = pd.DataFrame(x_selected_fea)
df_hdbscan_final_sub['CLUSTER'] = cluster_labels.reshape((-1, 1))

df_q1_res = pd.read_csv('./../data/kpro final/kpro-alldata-submission02.csv')
df_q1_res = df_q1_res.drop('Unnamed: 0', 1)
# add additional infos
df_q1_res.loc[df_q1_res['CLUSTER'] == 4, 'CLUSTER_CLASSIFIED'] = 'Level One'
df_q1_res.loc[df_q1_res['CLUSTER'] == 4, 'CLUSTER_SCORE'] = 5
df_q1_res.loc[df_q1_res['CLUSTER'] == 0, 'CLUSTER_CLASSIFIED'] = 'Level Two'
df_q1_res.loc[df_q1_res['CLUSTER'] == 0, 'CLUSTER_SCORE'] = 4
df_q1_res.loc[df_q1_res['CLUSTER'] == 2, 'CLUSTER_CLASSIFIED'] = 'Level Three'
df_q1_res.loc[df_q1_res['CLUSTER'] == 2, 'CLUSTER_SCORE'] = 3
df_q1_res.loc[df_q1_res['CLUSTER'] == 1, 'CLUSTER_CLASSIFIED'] = 'Level Four'
df_q1_res.loc[df_q1_res['CLUSTER'] == 1, 'CLUSTER_SCORE'] = 2
df_q1_res.loc[df_q1_res['CLUSTER'] == 3, 'CLUSTER_CLASSIFIED'] = 'Level Five'
df_q1_res.loc[df_q1_res['CLUSTER'] == 3, 'CLUSTER_SCORE'] = 1

df_q1_res_tar_slice = df_q1_res.loc[df_hdbscan_final_sub.index, :]

# Agg all data, combine them
df_q1_res_tar_slice['CODE_NAME'] = cluster_labels.reshape((-1, 1))
df_q1_res_tar_slice['CODE_NAME_POSSIBILITY'] = clusterer.probabilities_.reshape((-1, 1))
df_q1_res_tar_slice.to_csv('./../data/dbscan final/dbscan-slice-15_16-codename-submission01.csv')

df_hdbscan_final_sub['POSSIBILITY'] = clusterer.probabilities_.reshape((-1, 1))
df_hdbscan_final_sub['eventid'] = df_q1_res_tar_slice['eventid']
df_hdbscan_final_sub.to_csv('./../data/dbscan final/dbscan-ori-trainwithres-submission00.csv')

# calculate danger_confidence
x = pd.Series(df_q1_res_tar_slice.groupby(by=['CODE_NAME']).sum()['CLUSTER_SCORE'])
x = x.sort_values(ascending=False)
x.to_csv('./../data/dbscan final/dbscan-dangerlevel-sum-codename-submission02.csv')

# DOING distance ....
all_five_codename = x.index[1:6].tolist()
exemplars_ = clusterer.exemplars_
best_top_5_in_rank5 = np.asarray([exemplars_[i][0] for i in all_five_codename])

# 2017 needed to predicts
df_predict_index_target = [
    201701090031,
    201702210037,
    201703120023,
    201705050009,
    201705050010,
    201707010028,
    201707020006,
    201708110018,
    201711010006,
    201712010003
]

df_predict_from_source = df_target[df_target['eventid'].isin(df_predict_index_target)]

df_predict_selected_fea = df_predict_from_source[all_tar_idxs]
# x_selected_fea = df_dbscan_train_real[all_tar_idxs]

# encode provstate
le_pr_provstate = preprocessing.LabelEncoder()
en_pr_provstate = le_pr_provstate.fit_transform(df_predict_selected_fea['provstate'])
df_predict_selected_fea['provstate'] = en_pr_provstate.reshape((-1, 1))

# encode city
le_pr_city = preprocessing.LabelEncoder()
en_pr_city = le_pr_city.fit_transform(df_predict_selected_fea['city'])
df_predict_selected_fea['city'] = en_pr_city.reshape((-1, 1))

def calEuclideanDistance(vec1, vec2):
    dist = np.sqrt(np.sum(np.square(vec1 - vec2)))
    return dist

# calculate distance
all_dists = np.zeros((len(df_predict_index_target), len(all_five_codename)))
for i in range(len(df_predict_index_target)):
    for j in range(len(all_five_codename)):
        dist = calEuclideanDistance(df_predict_selected_fea.values[i], best_top_5_in_rank5[j])
        all_dists[i][j] = dist

print(all_dists)