import numpy as np
import pybullet as p
import pybullet_utils as pu
import grasp_utils as gu
import time
import pyquaternion
from collections import namedtuple
import numpy as np
import pybullet_utils as pu
from rrt import rrt
from rrt_connect import birrt
from itertools import combinations, product
import time

Grasp = namedtuple('Grasp', ['grasp_pose', 'grasp_jv', 'pre_grasp_pose', 'pre_grasp_jv',
                             'pos_distance', 'quat_distance', 'pre_pos_distance', 'pre_quat_distance'])

class UR5RobotiqPybulletController(object):
    # movable joints for each moveit group
    GROUPS = {
        'arm': ['shoulder_pan_joint', 'shoulder_lift_joint', 'elbow_joint',
                'wrist_1_joint', 'wrist_2_joint', 'wrist_3_joint'],
        'gripper': ['finger_joint', 'left_inner_knuckle_joint', 'left_inner_finger_joint', 'right_outer_knuckle_joint',
                    'right_inner_knuckle_joint', 'right_inner_finger_joint']
    }

    GROUP_INDEX = {
        'arm': [1, 2, 3, 4, 5, 6],
        'gripper': [9, 11, 13, 14, 16, 18]
    }

    LINK_COUNT = 10

    HOME = [0, -0.8227210029571718, -0.130, -0.660, 0, 1.62]
    RESET = [-np.pi, -np.pi / 2, np.pi / 2, -np.pi / 2, -np.pi / 2, 0]
    LIE = [0, 0, 0, 0, 0, 0]
    UP = [0, -1.5707, 0, -1.5707, 0, 0]
    OPEN_POSITION = [0] * 6
    CLOSED_POSITION = 0.4 * np.array([1, 1, -1, 1, 1, -1])
    EE_LINK_NAME = 'ee_link'
    TIP_LINK = "ee_link"
    BASE_LINK = "base_link"
    ARM = "arm"
    GRIPPER = "gripper"

    # this is read from moveit_configs joint_limits.yaml
    MOVEIT_ARM_MAX_VELOCITY = [3.15, 3.15, 3.15, 3.15, 3.15, 3.15]

    def __init__(self,
                 urdf='assets/ur5/ur5_robotiq.urdf',
                 base_pose=((0, 0, 0), (0, 0, 0, 1)),
                 initial_arm_joint_values=None):
        self.urdf = urdf
        self.base_pose = base_pose
        self.initial_arm_joint_values = initial_arm_joint_values if initial_arm_joint_values else self.RESET

        self.id = p.loadURDF(urdf,
                             basePosition=base_pose[0],
                             baseOrientation=base_pose[1],
                             useFixedBase=True,
                             flags=p.URDF_USE_SELF_COLLISION)

        joint_infos = [p.getJointInfo(self.id, joint_index)
                       for joint_index in range(p.getNumJoints(self.id))]
        self.JOINT_INDICES_DICT = {entry[1]: entry[0] for entry in joint_infos}
        self.EEF_LINK_INDEX = pu.link_from_name(self.id, self.EE_LINK_NAME)

        # for motion planning
        self.arm_difference_fn = pu.get_difference_fn(
            self.id, self.GROUP_INDEX['arm'])
        self.arm_distance_fn = pu.get_distance_fn(
            self.id, self.GROUP_INDEX['arm'])
        self.arm_sample_fn = pu.get_sample_fn(
            self.id, self.GROUP_INDEX['arm'])
        self.arm_extend_fn = pu.get_extend_fn(
            self.id, self.GROUP_INDEX['arm'])
        self.link_pairs = pu.get_self_link_pairs(
            self.id,
            self.GROUP_INDEX['arm'])
        self.closest_points_to_others = []
        self.closest_points_to_self = []
        self.max_distance_from_others = 0.5
        self.arm_difference_fn = pu.get_difference_fn(
            self.id, self.GROUP_INDEX['arm'])
        self.attach_cid = None
        self.attach_object_id = None
        self.reset()

    def reset(self):
        self.set_arm_joints(self.initial_arm_joint_values)
        self.set_gripper_joints(self.OPEN_POSITION)
        if self.attach_cid is not None:
            p.removeConstraint(self.attach_cid)
        self.attach_cid = None
        self.attach_object_id = None

    def set_arm_joints(self, joint_values):
        pu.set_joint_positions(self.id, self.GROUP_INDEX['arm'], joint_values)
        pu.control_joints(self.id, self.GROUP_INDEX['arm'], joint_values)

    def control_arm_joints(self, joint_values):
        pu.control_joints(
            body=self.id,
            joints=self.GROUP_INDEX['arm'],
            positions=joint_values)

    def set_gripper_joints(self, joint_values):
        pu.set_joint_positions(
            self.id, self.GROUP_INDEX['gripper'], joint_values)
        pu.control_joints(self.id, self.GROUP_INDEX['gripper'], joint_values)

    def control_gripper_joints(self, joint_values):
        pu.control_joints(self.id, self.GROUP_INDEX['gripper'], joint_values)

    def close_gripper(self, realtime=False, duration=2):
        waypoints = self.plan_gripper_joint_values(
            self.CLOSED_POSITION, duration=duration)
        self.execute_gripper_plan(waypoints, realtime)

    def open_gripper(self, realtime=False, duration=2):
        waypoints = self.plan_gripper_joint_values(
            self.OPEN_POSITION, duration=duration)
        self.execute_gripper_plan(waypoints, realtime)

    def plan_gripper_joint_values(self, goal_joint_values, start_joint_values=None, duration=1):
        if start_joint_values is None:
            start_joint_values = self.get_gripper_joint_values()
        num_steps = int(duration * 240)
        discretized_plan = np.linspace(
            start_joint_values, goal_joint_values, num_steps)
        # same as plan arm joint values
        discretized_plan = np.vstack((discretized_plan, np.repeat(discretized_plan[-1, None], 240, axis=0)))
        return discretized_plan

    def get_arm_fk_pybullet(self, joint_values):
        return pu.forward_kinematics(self.id, self.GROUP_INDEX['arm'], joint_values, self.EEF_LINK_INDEX)

    def get_end_effector_pose(self):
        return self.get_arm_fk_pybullet(self.get_arm_joint_values())

    def get_arm_ik_pybullet(self, pose_2d, arm_joint_values=None, gripper_joint_values=None):
        # This get IK does not check collision (self and with others)
        # gripper_joint_values = self.get_gripper_joint_values() if gripper_joint_values is None else gripper_joint_values
        # arm_joint_values = self.get_arm_joint_values() if arm_joint_values is None else arm_joint_values

        joint_values = p.calculateInverseKinematics(self.id,
                                                    # self.JOINT_INDICES_DICT[self.EEF_LINK],
                                                    self.EEF_LINK_INDEX,
                                                    pose_2d[0],
                                                    pose_2d[1],
                                                    maxNumIterations=10000,
                                                    residualThreshold=.0001)
        ik_result = list(joint_values[:6])
        actual_pose = self.get_arm_fk_pybullet(ik_result)
        pos_distance = np.linalg.norm(
            np.array(actual_pose[0]) - np.array(pose_2d[0]))
        quat_distance = pyquaternion.Quaternion.absolute_distance(
            pyquaternion.Quaternion(pu.change_quat_rep(pose_2d[1])),
            pyquaternion.Quaternion(pu.change_quat_rep(actual_pose[1])))
        return ik_result, pos_distance, quat_distance

    def plan_arm_joint_values_simple(self, goal_joint_values, start_joint_values=None, duration=None):
        """ Linear interpolation between joint_values """
        start_joint_values = self.get_arm_joint_values(
        ) if start_joint_values is None else start_joint_values

        diffs = self.arm_difference_fn(goal_joint_values, start_joint_values)
        # specifically handle the wrist joint, whose period should be pi instead of 2*pi
        diffs[-1] = (diffs[-1] + np.pi / 2) % np.pi - np.pi / 2
        steps = np.abs(np.divide(diffs, self.MOVEIT_ARM_MAX_VELOCITY)) * 240
        num_steps = int(max(steps))
        if duration is not None:
            num_steps = int(duration * 240)
            # num_steps = max(int(duration * 240), steps)     # this should ensure that it satisfies the max velocity of the end-effector

        goal_joint_values = np.array(start_joint_values) + np.array(diffs)
        waypoints = np.linspace(
            start_joint_values, goal_joint_values, num_steps)

        # duplicate the last waypoint for 240 times to make sure the target is reached
        # because the task executer steps every task for 1/240 but the low level controller simulates another 1s in the end
        waypoints = np.vstack((waypoints, np.repeat(waypoints[-1, None], 240, axis=0)))
        return waypoints

    @staticmethod
    def convert_range(joint_values):
        """ Convert continuous joint to range [-pi, pi] """
        circular_idx = [0, 3, 4, 5]
        new_joint_values = []
        for i, v in enumerate(joint_values):
            if i in circular_idx:
                new_joint_values.append(pu.wrap_angle(v))
            else:
                new_joint_values.append(v)
        return new_joint_values

    # execution
    def execute_arm_plan(self, plan, realtime=False):
        """
        execute a discretized arm plan (list of waypoints)
        """
        for wp in plan:
            self.control_arm_joints(wp)
            p.stepSimulation()
            if realtime:
                time.sleep(1. / 240.)
        pu.step(1)

    def execute_gripper_plan(self, plan, realtime=False):
        """
        execute a discretized gripper plan (list of waypoints)
        """
        for wp in plan:
            self.control_gripper_joints(wp)
            p.stepSimulation()
            if realtime:
                time.sleep(1. / 240.)
        pu.step(1)

    def step(self):
        """ step the robot for 1/240 second """
        # calculate the latest conf and control array
        if self.arm_discretized_plan is None or self.arm_wp_target_index == len(self.arm_discretized_plan):
            pass
        else:
            self.control_arm_joints(
                self.arm_discretized_plan[self.arm_wp_target_index])
            self.arm_wp_target_index += 1

        if self.gripper_discretized_plan is None or self.gripper_wp_target_index == len(self.gripper_discretized_plan):
            pass
        else:
            self.control_gripper_joints(
                self.gripper_discretized_plan[self.gripper_wp_target_index])
            self.gripper_wp_target_index += 1

    def equal_conf(self, conf1, conf2, tol=0):
        adapted_conf2 = self.adapt_conf(conf2, conf1)
        return np.allclose(conf1, adapted_conf2, atol=tol)

    def adapt_conf(self, conf2, conf1):
        """ adapt configuration 2 to configuration 1"""
        diff = self.arm_difference_fn(conf2, conf1)
        adapted_conf2 = np.array(conf1) + np.array(diff)
        return adapted_conf2.tolist()

    def get_arm_joint_values(self):
        return pu.get_joint_positions(self.id, self.GROUP_INDEX['arm'])

    def get_gripper_joint_values(self):
        return pu.get_joint_positions(self.id, self.GROUP_INDEX['gripper'])

    def get_eef_pose(self):
        return pu.get_link_pose(self.id, self.EEF_LINK_INDEX)

    def update_closest_points(self, excluded_objects=[]):
        others_id = [p.getBodyUniqueId(i)
                     for i in range(p.getNumBodies())
                     if p.getBodyUniqueId(i) != self.id]
        for i in excluded_objects:
            others_id.remove(i)
        self.closest_points_to_others = [
            sorted(list(p.getClosestPoints(
                bodyA=self.id, bodyB=other_id,
                distance=self.max_distance_from_others)),
                key=lambda contact_points: contact_points[8])
            for other_id in others_id]
        self.closest_points_to_self = [
            p.getClosestPoints(
                bodyA=self.id, bodyB=self.id,
                distance=0,
                linkIndexA=link1, linkIndexB=link2)
            for link1, link2 in self.link_pairs]

    def check_collision(self, conf=None, collision_distance=0.0, excluded_objects=[], self_collision=True):
        if conf is not None:
            self.set_arm_joints(conf)
        self.update_closest_points(excluded_objects)
        # Collisions with others
        for i, closest_points_to_other in enumerate(
                self.closest_points_to_others):
            # Check whether closest point's distance is
            # less than collision distance
            for point in closest_points_to_other:
                if point[8] < collision_distance:
                    return True, (point[1], point[2], point[3], point[4])
        if self_collision:
            # Self Collision
            for closest_points_to_self_link in self.closest_points_to_self:
                if len(closest_points_to_self_link) > 0:
                    point = closest_points_to_self_link[0]
                    return True, (point[1], point[2], point[3], point[4])
        return False, None

    def plan_grasp(self, target):
        current_jv = self.get_arm_joint_values()
        grasps = []
        for g, p_g in zip(target.grasps_eef, target.pre_grasps_eef):
            g_pose = pu.split_7d(g)
            pre_g_pose = pu.split_7d(p_g)
            g_pose = gu.convert_grasp_in_object_to_world(
                target.get_pose(), g_pose)
            pre_g_pose = gu.convert_grasp_in_object_to_world(
                target.get_pose(), pre_g_pose)
            jv, pos_distance, quat_distance = self.get_arm_ik_pybullet(g_pose)
            pre_jv, pre_pos_distance, pre_quat_distance = self.get_arm_ik_pybullet(
                pre_g_pose)
            self.set_arm_joints(jv)
            # exclude target for collision checking
            grasp_collision, info = self.check_collision(
                excluded_objects=[target.id])
            self.set_arm_joints(pre_jv)
            pre_grasp_collision, info = self.check_collision()
            collision = pre_grasp_collision or grasp_collision
            if not collision and \
                    pos_distance < 1e-3 and \
                    quat_distance < 1e-3 and \
                    pre_pos_distance < 1e-3 and \
                    pre_quat_distance < 1e-3:
                grasps.append(Grasp(grasp_pose=g_pose,
                                    grasp_jv=jv,
                                    pre_grasp_pose=pre_g_pose,
                                    pre_grasp_jv=pre_jv,
                                    pos_distance=pos_distance,
                                    quat_distance=quat_distance,
                                    pre_pos_distance=pre_pos_distance,
                                    pre_quat_distance=pre_quat_distance))
        self.set_arm_joints(current_jv)
        return grasps

    def cartesian_control(self, axis, simple=True, value=0.1, realtime=False, excluded_objects=[], duration=2):
        target_pose = self.get_eef_pose()
        if axis == 'x':
            i = 0
        elif axis == 'y':
            i = 1
        elif axis == 'z':
            i = 2
        else:
            raise TypeError('unsupported axis')
        target_pose[0][i] += value
        target_jv = self.get_arm_ik_pybullet(target_pose)[0]
        plan = self.plan_arm_joint_values_simple(target_jv, duration=duration) \
            if simple \
            else self.plan_arm_joint_values(target_jv,
                                            excluded_objects=excluded_objects)
        self.execute_arm_plan(plan, realtime)
        return plan

    def move_arm_to(self,
                    target_jv,
                    simple=True,
                    realtime=False,
                    excluded_objects=[],
                    duration=2):
        if simple:
            plan = self.plan_arm_joint_values_simple(target_jv, duration=duration)
        else:
            plan = self.plan_arm_joint_values(target_jv, excluded_objects)
        self.execute_arm_plan(plan, realtime=realtime)
        return plan

    def plan_arm_joint_values(self,
                              goal_conf,
                              excluded_objects=[],
                              start_conf=None,
                              planner='birrt',
                              smooth=200,
                              greedy=True,
                              goal_tolerance=0.001,
                              goal_bias=0.2,
                              resolutions=0.05,
                              iterations=2000,
                              restarts=10):
        current_conf = self.get_arm_joint_values()
        start_conf = current_conf if start_conf is None else start_conf

        # get some functions
        def collision_fn(conf):
            return self.check_collision(
                conf, excluded_objects=excluded_objects)[0]

        goal_test = pu.get_goal_test_fn(goal_conf, goal_tolerance)
        extend_fn = self.arm_extend_fn

        if planner == 'rrt':
            for i in range(restarts):
                iter_start = time.time()
                path_conf = rrt(start=start_conf,
                                goal_sample=goal_conf,
                                distance=self.arm_difference_fn,
                                sample=self.arm_sample_fn,
                                extend=extend_fn,
                                collision=collision_fn,
                                goal_probability=goal_bias,
                                iterations=iterations,
                                goal_test=goal_test,
                                greedy=greedy,
                                visualize=True,
                                fk=self.get_arm_fk_pybullet,
                                group=False)
                iter_time = time.time() - iter_start
                self.set_arm_joints(current_conf)
                if path_conf is None:
                    print('trial {} ({} iterations) fails in {:.2f} seconds'.format(
                        i + 1, iterations, iter_time))
                    pu.remove_all_markers()
                else:
                    return path_conf
        elif planner == 'birrt':
            for i in range(restarts):
                iter_start = time.time()
                path_conf = birrt(start_conf=start_conf,
                                  goal_conf=goal_conf,
                                  distance=self.arm_difference_fn,
                                  sample=self.arm_sample_fn,
                                  extend=extend_fn,
                                  collision=collision_fn,
                                  iterations=iterations,
                                  smooth=smooth,
                                  visualize=True,
                                  fk=self.get_arm_fk_pybullet,
                                  group=False,
                                  greedy=greedy)
                iter_time = time.time() - iter_start
                self.set_arm_joints(current_conf)
                if path_conf is None:
                    print('trial {} ({} iterations) fails in {:.2f} seconds'.format(
                        i + 1, iterations, iter_time))
                    pu.remove_all_markers()
                else:
                    return path_conf
        else:
            raise ValueError('planner must be in \'rrt\' or \'birrt\'')

    def attach_object(self, target_id):
        target_pose = pu.get_body_pose(target_id)
        eef_pose = self.get_eef_pose()
        eef_P_world = p.invertTransform(eef_pose[0], eef_pose[1])
        eef_P_target = p.multiplyTransforms(
            eef_P_world[0], eef_P_world[1], target_pose[0], target_pose[1])
        self.attach_cid = p.createConstraint(
            parentBodyUniqueId=target_id,
            parentLinkIndex=-1,
            childBodyUniqueId=self.id,
            childLinkIndex=self.EEF_LINK_INDEX,
            jointType=p.JOINT_FIXED,
            jointAxis=[0, 0, 0],
            parentFramePosition=[0, 0, 0],
            childFramePosition=eef_P_target[0],
            childFrameOrientation=eef_P_target[1])
        self.attach_object_id = target_id

    def detach(self):
        p.removeConstraint(self.attach_cid)
        self.attach_cid = None
        self.attach_object_id = None

    def global_to_ur5_frame(self, position, rotation=None):
        self_pos, self_rot = p.getBasePositionAndOrientation(self.id)
        invert_self_pos, invert_self_rot = p.invertTransform(
            self_pos, self_rot)
        ur5_frame_pos, ur5_frame_rot = p.multiplyTransforms(
            invert_self_pos, invert_self_rot,
            position, invert_self_rot if rotation is None else rotation
        )
        return ur5_frame_pos, ur5_frame_rot

    def get_link_global_positions(self):
        linkstates = [p.getLinkState(
            self.id, link_id, computeForwardKinematics=True)
            for link_id in range(UR5RobotiqPybulletController.LINK_COUNT)]
        link_world_positions = [
            world_pos for
            world_pos, world_rot, _, _, _, _, in linkstates
        ]
        return link_world_positions

    def get_pose(self):
        return p.getBasePositionAndOrientation(self.id)


import os
import tensorflow as tf
from keras.preprocessing import image
from flask import Flask, request, render_template, send_from_directory
from keras.models import load_model
from keras.applications.resnet50 import preprocess_input, decode_predictions, ResNet50
import cv2
import numpy as np

app = Flask(__name__)

APP_ROOT = os.path.dirname(os.path.abspath(__file__))

from extract_bottleneck_features import *
dognames_txt = open("dognames.txt","r")
dog_names     = dognames_txt.read().split('\n')
Xception_model = load_model('saved_models/Xception.best_weights.hdf5')

# define ResNet50 model
ResNet50_model = ResNet50(weights='imagenet')

def path_to_tensor(img_path):
    # loads RGB image as PIL.Image.Image type
    img = image.load_img(img_path, target_size=(224, 224))
    # convert PIL.Image.Image type to 3D tensor with shape (224, 224, 3)
    x = image.img_to_array(img)
    # convert 3D tensor to 4D tensor with shape (1, 224, 224, 3) and return 4D tensor
    return np.expand_dims(x, axis=0)

def Xception_predictbreed(img_path):
    Xception_bottlenecks = extract_Xception(path_to_tensor(img_path))
    Xception_prediction  = Xception_model.predict(Xception_bottlenecks)
    return dog_names[np.argmax(Xception_prediction)]

def ResNet50_predict_labels(img_path):
    # returns prediction vector for image located at img_path
    img = preprocess_input(path_to_tensor(img_path))
    return np.argmax(ResNet50_model.predict(img))

### returns "True" if a dog is detected in the image stored at img_path
def dog_detector(img_path):
    prediction = ResNet50_predict_labels(img_path)
    return ((prediction <= 268) & (prediction >= 151))

# extract pre-trained face detector
face_cascade = cv2.CascadeClassifier('haarcascades/haarcascade_frontalface_alt.xml')

def face_detector(img_path):
    img = cv2.imread(img_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray)
    return len(faces) > 0

def Xception_detector(img_path):
    fd = face_detector(img_path)         #detect human face
    dd = dog_detector(img_path)          #detect dog
    pd = Xception_predictbreed(img_path) #predicted breed

    if(fd==True and dd==False):
        s1 = "You are human, and the dog breed you resemble most is "
    elif(fd==False and dd==True):
        if(pd[0]=="A" or pd[0]=="E" or pd[0]=="I" or pd[0]=="O" or pd[0]=="U"):
            s1 = "You are an "
        else:
            s1 = "You are a "
    else:
        s1 = "Well, the dog breed you resemble most is "
    return s1 + pd

@app.route("/")
def index():
    return render_template("upload.html")

@app.route("/upload", methods=["POST"])
def upload():
    target = os.path.join(APP_ROOT, 'images/')
    print(target)
    if not os.path.isdir(target):
            os.mkdir(target)
    else:
        print("Upload directory: {}".format(target), "already exists")
    filename = ""
    destination = ""
    detection = " "
    for upload in request.files.getlist("file"):
        print(upload)
        print("{} is the file name".format(upload.filename))
        filename = upload.filename
        destination = "/".join([target, filename])
        print ("Save image file to:", destination)
        upload.save(destination)
        detection = Xception_detector(destination)
    return render_template("complete.html", image_name=filename, text_str = detection)

@app.route('/upload/<filename>')
def send_image(filename):
    return send_from_directory("images", filename)


import argparse
import os
import random
from math import log10
import scipy.stats as stats
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import torchvision.utils as vutils
from tensorboardX import SummaryWriter
from torch.autograd import Variable, grad
from models.dev_model import *
from data.nvData import CreateDataLoader

parser = argparse.ArgumentParser()
parser.add_argument('--datarootC', required=True, help='path to colored dataset')
parser.add_argument('--datarootS', required=True, help='path to sketch dataset')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=4)
parser.add_argument('--batchSize', type=int, default=16, help='input batch size')
parser.add_argument('--imageSize', type=int, default=256, help='the height / width of the input image to network')
parser.add_argument('--cut', type=int, default=1, help='cut backup frequency')
parser.add_argument('--niter', type=int, default=700, help='number of epochs to train for')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ndf', type=int, default=64)
parser.add_argument('--lrG', type=float, default=0.0001, help='learning rate, default=0.0001')
parser.add_argument('--lrD', type=float, default=0.0001, help='learning rate, default=0.0001')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--cuda', action='store_true', help='enables cuda')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--optim', action='store_true', help='load optimizer\'s checkpoint')
parser.add_argument('--outf', default='.', help='folder to output images and model checkpoints')
parser.add_argument('--Diters', type=int, default=1, help='number of D iters per each G iter')
parser.add_argument('--manualSeed', type=int, default=2345, help='random seed to use. Default=1234')
parser.add_argument('--baseGeni', type=int, default=2500, help='start base of pure pair L1 loss')
parser.add_argument('--geni', type=int, default=0, help='continue gen image num')
parser.add_argument('--epoi', type=int, default=0, help='continue epoch num')
parser.add_argument('--env', type=str, default=None, help='tensorboard env')
parser.add_argument('--advW', type=float, default=0.0001, help='adversarial weight, default=0.0001')
parser.add_argument('--gpW', type=float, default=10, help='gradient penalty weight')

opt = parser.parse_args()
print(opt)

####### regular set up
if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")
gen_iterations = opt.geni
try:
    os.makedirs(opt.outf)
except OSError:
    pass
# random seed setup                                  # !!!!!
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
    torch.cuda.manual_seed(opt.manualSeed)
cudnn.benchmark = True
####### regular set up end

writer = SummaryWriter(log_dir=opt.env, comment='this is great')

dataloader = CreateDataLoader(opt)

netG = def_netG(ngf=opt.ngf)
if opt.netG != '':
    netG.load_state_dict(torch.load(opt.netG))
print(netG)

netD = def_netD(ndf=opt.ndf)
if opt.netD != '':
    netD.load_state_dict(torch.load(opt.netD))
print(netD)

netF = def_netF()
print(netD)

criterion_L1 = nn.L1Loss()
criterion_L2 = nn.MSELoss()
one = torch.FloatTensor([1])
mone = one * -1

fixed_sketch = torch.FloatTensor()
fixed_hint = torch.FloatTensor()
saber = torch.FloatTensor([0.485 - 0.5, 0.456 - 0.5, 0.406 - 0.5]).view(1, 3, 1, 1)
diver = torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)

if opt.cuda:
    netD.cuda()
    netG.cuda()
    netF.cuda()
    fixed_sketch, fixed_hint = fixed_sketch.cuda(), fixed_hint.cuda()
    saber, diver = saber.cuda(), diver.cuda()
    criterion_L1.cuda()
    criterion_L2.cuda()
    one, mone = one.cuda(), mone.cuda()

# setup optimizer
optimizerG = optim.Adam(netG.parameters(), lr=opt.lrG, betas=(opt.beta1, 0.9))
optimizerD = optim.Adam(netD.parameters(), lr=opt.lrD, betas=(opt.beta1, 0.9))

if opt.optim:
    optimizerG.load_state_dict(torch.load('%s/optimG_checkpoint.pth' % opt.outf))
    optimizerD.load_state_dict(torch.load('%s/optimD_checkpoint.pth' % opt.outf))

# schedulerG = lr_scheduler.ReduceLROnPlateau(optimizerG, mode='max', verbose=True, min_lr=0.0000005,
#                                             patience=8)  # 1.5*10^5 iter
# schedulerD = lr_scheduler.ReduceLROnPlateau(optimizerD, mode='max', verbose=True, min_lr=0.0000005,
#                                             patience=8)  # 1.5*10^5 iter

# schedulerG = lr_scheduler.MultiStepLR(optimizerG, milestones=[60, 120], gamma=0.1)  # 1.5*10^5 iter
# schedulerD = lr_scheduler.MultiStepLR(optimizerD, milestones=[60, 120], gamma=0.1)

def calc_gradient_penalty(netD, real_data, fake_data):
    # print "real_data: ", real_data.size(), fake_data.size()
    alpha = torch.rand(opt.batchSize, 1, 1, 1)
    # alpha = alpha.expand(opt.batchSize, real_data.nelement() / opt.batchSize).contiguous().view(opt.batchSize, 3, 64,
    #                                                                                64)
    alpha = alpha.cuda() if opt.cuda else alpha

    interpolates = alpha * real_data + ((1 - alpha) * fake_data)

    if opt.cuda:
        interpolates = interpolates.cuda()
    interpolates = Variable(interpolates, requires_grad=True)

    disc_interpolates = netD(interpolates)

    gradients = grad(outputs=disc_interpolates, inputs=interpolates,
                     grad_outputs=torch.ones(disc_interpolates.size()).cuda() if opt.cuda else torch.ones(
                         disc_interpolates.size()),
                     create_graph=True, retain_graph=True, only_inputs=True)[0]

    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * opt.gpW
    return gradient_penalty

flag = 1
lower, upper = 0, 1
mu, sigma = 1, 0.01
maskS = opt.imageSize // 4
X = stats.truncnorm(
    (lower - mu) / sigma, (upper - mu) / sigma, loc=mu, scale=sigma)
for epoch in range(opt.niter):
    data_iter = iter(dataloader)
    i = 0
    while i < len(dataloader):
        ############################
        # (1) Update D network
        ###########################
        for p in netD.parameters():  # reset requires_grad
            p.requires_grad = True  # they are set to False below in netG update
        for p in netG.parameters():
            p.requires_grad = False  # to avoid computation

        # train the discriminator Diters times
        Diters = opt.Diters

        if gen_iterations < opt.baseGeni:  # L1 stage
            Diters = 0

        j = 0
        while j < Diters and i < len(dataloader):

            j += 1
            netD.zero_grad()

            data = data_iter.next()
            real_cim, real_vim, real_sim = data
            i += 1
            ###############################

            if opt.cuda:
                real_cim, real_vim, real_sim = real_cim.cuda(), real_vim.cuda(), real_sim.cuda()

            mask = torch.cat([torch.rand(1, 1, maskS, maskS).ge(X.rvs(1)[0]).float() for _ in range(opt.batchSize)],
                             0).cuda()
            hint = torch.cat((real_vim * mask, mask), 1)

            # train with fake

            fake_cim = netG(Variable(real_sim, volatile=True), Variable(hint, volatile=True)).data
            errD_fake = netD(Variable(torch.cat((fake_cim, real_sim), 1))).mean(0).view(1)
            errD_fake.backward(one, retain_graph=True)  # backward on score on real

            errD_real = netD(Variable(torch.cat((real_cim, real_sim), 1))).mean(0).view(1)
            errD_real.backward(mone, retain_graph=True)  # backward on score on real

            errD = errD_real - errD_fake

            # gradient penalty
            gradient_penalty = calc_gradient_penalty(netD, torch.cat([real_cim, real_sim], 1),
                                                     torch.cat([fake_cim, real_sim], 1))
            gradient_penalty.backward()
            optimizerD.step()

        ############################
        # (2) Update G network
        ############################
        if i < len(dataloader):
            for p in netD.parameters():
                p.requires_grad = False  # to avoid computation
            for p in netG.parameters():
                p.requires_grad = True  # to avoid computation
            netG.zero_grad()

            data = data_iter.next()
            real_cim, real_vim, real_sim = data
            i += 1

            if opt.cuda:
                real_cim, real_vim, real_sim = real_cim.cuda(), real_vim.cuda(), real_sim.cuda()

            mask = torch.cat([torch.rand(1, 1, maskS, maskS).ge(X.rvs(1)[0]).float() for _ in range(opt.batchSize)],
                             0).cuda()
            hint = torch.cat((real_vim * mask, mask), 1)

            if flag:  # fix samples
                writer.add_image('target imgs', vutils.make_grid(real_cim.mul(0.5).add(0.5), nrow=16))
                writer.add_image('sketch imgs', vutils.make_grid(real_sim.mul(0.5).add(0.5), nrow=16))
                writer.add_image('hint', vutils.make_grid((real_vim * mask).mul(0.5).add(0.5), nrow=16))
                vutils.save_image(real_cim.mul(0.5).add(0.5),
                                  '%s/color_samples' % opt.outf + '.png')
                vutils.save_image(real_sim.mul(0.5).add(0.5),
                                  '%s/blur_samples' % opt.outf + '.png')
                fixed_sketch.resize_as_(real_sim).copy_(real_sim)
                fixed_hint.resize_as_(hint).copy_(hint)

                flag -= 1

            fake = netG(Variable(real_sim), Variable(hint))

            if gen_iterations < opt.baseGeni:
                contentLoss = criterion_L2(netF((fake.mul(0.5) - Variable(saber)) / Variable(diver)),
                                           netF(Variable((real_cim.mul(0.5) - saber) / diver)))
                contentLoss.backward()
                errG = contentLoss
                # contentLoss = criterion_L1(fake, Variable(real_cim))
                # contentLoss.backward()
                # errG = contentLoss
            else:
                errG = netD(torch.cat((fake, Variable(real_sim)), 1)).mean(0).view(
                    1) * opt.advW  # TODO: what if???
                errG.backward(mone, retain_graph=True)

                contentLoss = criterion_L2(netF((fake.mul(0.5) - Variable(saber)) / Variable(diver)),
                                           netF(Variable((real_cim.mul(0.5) - saber) / diver)))
                contentLoss.backward()
                # contentLoss = criterion_L1(fake, Variable(real_cim))
                # contentLoss.backward(retain_graph=True)

            optimizerG.step()

        ############################
        # (3) Report & 100 Batch checkpoint
        ############################

        if gen_iterations < opt.baseGeni:
            writer.add_scalar('VGG MSE Loss', contentLoss.data[0], gen_iterations)
            print('[%d/%d][%d/%d][%d] content %f '
                  % (epoch, opt.niter, i, len(dataloader), gen_iterations, contentLoss.data[0]))
        else:
            writer.add_scalar('VGG MSE Loss', contentLoss.data[0], gen_iterations)
            writer.add_scalar('wasserstein distance', errD.data[0], gen_iterations)
            writer.add_scalar('errD_real', errD_real.data[0], gen_iterations)
            writer.add_scalar('errD_fake', errD_fake.data[0], gen_iterations)
            writer.add_scalar('Gnet loss toward real', errG.data[0], gen_iterations)
            writer.add_scalar('gradient_penalty', gradient_penalty.data[0], gen_iterations)
            print('[%d/%d][%d/%d][%d] errD: %f err_G: %f err_D_real: %f err_D_fake %f content loss %f'
                  % (epoch, opt.niter, i, len(dataloader), gen_iterations,
                     errD.data[0], errG.data[0], errD_real.data[0], errD_fake.data[0], contentLoss.data[0]))

        if gen_iterations % 500 == 0:
            fake = netG(Variable(fixed_sketch, volatile=True), Variable(fixed_hint, volatile=True))
            writer.add_image('colorized imgs', vutils.make_grid(fake.data.mul(0.5).add(0.5), nrow=16),
                             gen_iterations)

        if gen_iterations % 2000 == 0:
            for name, param in netG.named_parameters():
                writer.add_histogram('netG ' + name, param.clone().cpu().data.numpy(), gen_iterations)
            for name, param in netD.named_parameters():
                writer.add_histogram('netD ' + name, param.clone().cpu().data.numpy(), gen_iterations)
            vutils.save_image(fake.data.mul(0.5).add(0.5),
                              '%s/fake_samples_gen_iter_%08d.png' % (opt.outf, gen_iterations))
        gen_iterations += 1

    # do checkpointing
    if opt.cut == 0:
        torch.save(netG.state_dict(), '%s/netG_epoch_only.pth' % opt.outf)
        torch.save(netD.state_dict(), '%s/netD_epoch_only.pth' % opt.outf)
    elif epoch % opt.cut == 0:
        torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (opt.outf, epoch))
        torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (opt.outf, epoch))
    torch.save(optimizerG.state_dict(), '%s/optimG_checkpoint.pth' % opt.outf)
    torch.save(optimizerD.state_dict(), '%s/optimD_checkpoint.pth' % opt.outf)


import cv2
import math
import numpy as np
from retinaface import RetinaFace
from facealign import faceAlign

thresh = 0.4
gpuid = -1 #if gpu, set 0
resize_size = 640
save_size = 224
use_face_align = 1
prehead = "../"

detector = RetinaFace('./mnet.25/mnet.25', 0, gpuid, 'net3')

def resizeImg(img):
    m_resizeX, m_resizeY = 1.0, 1.0;
    H, W, C = img.shape

    m_resizeX = m_resizeY = max(H / resize_size, W / resize_size)

    img = cv2.resize(img, ((int)(W / m_resizeX), (int)(H / m_resizeY)))
    img = cv2.copyMakeBorder(img, 0, max(resize_size - H, 0), 0, max(resize_size - W, 0), 0)

    return img

def getMaxFace(img):
    detect_tab = 0
    box, landmark = cropFace(img)
    if len(box):
        detect_tab = 1

        if use_face_align:
            img = faceAlign(img, landmark)

        else:
            w = (box[3] - box[1]) + 1
            h = (box[2] - box[0]) + 1
            l = max(w, h)

            x0 = int(max(0, (int)(box[1] + box[3] - l) / 2))
            y0 = int(max(0, (int)(box[0] + box[2] - l) / 2))
            h = int(min(img.shape[0], l))
            w = int(min(img.shape[1], l))

            img = img[x0 : x0 + w, y0 : y0 + h]

    return img, detect_tab

def cropFace(img):
    im_shape = img.shape
    scales = [min(im_shape[0],im_shape[1]), max(im_shape[0],im_shape[1])]

    target_size = scales[0]
    max_size = scales[1]
    im_size_min = np.min(im_shape[0:2])
    im_size_max = np.max(im_shape[0:2])

    im_scale = float(target_size) / float(im_size_min)

    if np.round(im_scale * im_size_max) > max_size:
        im_scale = float(max_size) / float(im_size_max)

    scales = [im_scale]
    flip = False

    faces, landmarks = detector.detect(img, thresh, scales=scales, do_flip=flip)

    if len(faces):
        idx = 0
        for i in range(1,len(faces)):
            boxI = faces[i].astype(np.int)
            areaI = (boxI[3] - boxI[1]) * (boxI[2] - boxI[0])
            boxIdx = faces[idx].astype(np.int)
            areaIdx = (boxIdx[3] - boxIdx[1]) * (boxIdx[2] - boxIdx[0])
            if areaI > areaIdx:
               idx = i

        box = faces[idx].astype(np.int)
        landmark = landmarks[idx].astype(np.int)

        return box, landmark
    else:
        return [], []

def detectFace(load_path, save_path,
    save_noface_A = "../no_faces_A",
    txt_path = "../faceList.txt"):
    line_list = []
    with open(txt_path, "r") as f:
        line_list = f.readlines()
        for n in range(len(line_list)):
            line = "/" + line_list[n].replace("\n", "")
            if line:
                img = cv2.imread(prehead + load_path + line)

                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img, detect_tab = getMaxFace(img)
                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

                cv2.imwrite(prehead + save_path  + line, img)
                if not detect_tab:
                    cv2.imwrite(save_noface_A + line.replace("/", "#"), img)

            print("Process: "+ str(n + 1) + "/" + str(len(line_list)) + " " + line)

def detectFaceWithLabel(load_path, save_path, label_path, tag,
    save_noface_T = "../no_faces_T",
    save_noface_F = "../no_faces_F",
    txt_path = "../faceList.txt"):

    label_dict = {}
    line_list = []
    with open(prehead + label_path, "r") as f:
        line_list = f.readlines()
        for n in range(len(line_list)):
            line = line_list[n].replace("\n", "")
            if line:
                line = line.split(" ")
                key = line[0].replace(tag, "")
                value = line[1]
                label_dict.update({key : value})

    with open(txt_path, "r") as f:
        line_list = f.readlines()
        for n in range(len(line_list)):
            line = "/" + line_list[n].replace("\n", "")
            if line:
                img = cv2.imread(prehead + load_path + "/" + line)

                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img = resizeImg(img)
                img, detect_tab = getMaxFace(img)
                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

                if detect_tab:
                    cv2.imwrite(prehead + save_path + "/" + line, img)
                else:
                    if label_dict[line] == "1":
                        cv2.imwrite(save_noface_T + "/" + line.replace("/", "#"), img)
                    else:
                        cv2.imwrite(save_noface_F + "/" + line.replace("/", "#"), img)
            print("Process: "+ str(n + 1) + "/" + str(len(line_list)) + " " + line)

# detectFaceWithLabel("train", "save_merges_train", "train_label.txt", tag = "train")
detectFace("val", "save_merges_val")
# detectFaceWithLabel("val", "save_merges_val", "val_label.txt", tag = "val")
# detectFace("test", "save_merges_test")


from flask import current_app

from app import db
from app.network_x_tools.network_x_utils import network_x_utils
from app.myths.process_myths import process_myths
from app.errors.errors import CustomError
from app.models import Scores

import numpy as np
import random
import pickle
import sklearn  # This is used even though it isn't directly referenced
import sys

class process_solutions:

    """

    Solutions can be divided into two categories, mitigations and adaptations. Both of
    these are shown as potential solutions to climate issues affecting the user.

    Class used for solution related functions. Uses a fresh copy of the NetworkX graph to
    return relevant information about solutions.

    Any additional functions related to solutions should be added here.

    """

    def __init__(self, max_solutions=4, adaptation_to_mitigation_ratio=0.5):
        self.G = current_app.config["G"].copy()
        self.NX_UTILS = network_x_utils()
        self.MYTH_PROCESS = process_myths()
        self.MAX_SOLUTIONS = max_solutions
        self.A_TO_M_RATIO = adaptation_to_mitigation_ratio

    def solution_randomizer(self, adaptation_solutions, mitigation_solutions):
        """Takes list of solutions and decides which adaptation solutions to randomly
        show and which mitigation solutions to randomly show.

        Parameters
        ----------
        adaptation_solutions - A list of adaptation solutions
        mititgation_solutions - A list of mitigation solutions
        """
        number_adaptation_max = np.math.floor(self.MAX_SOLUTIONS * self.A_TO_M_RATIO)
        if len(adaptation_solutions) <= number_adaptation_max:
            number_mitigation = self.MAX_SOLUTIONS - len(adaptation_solutions)
            solutions = adaptation_solutions + random.sample(
                mitigation_solutions, number_mitigation
            )
        else:
            solutions = random.sample(
                adaptation_solutions, number_adaptation_max
            ) + random.sample(
                mitigation_solutions, self.MAX_SOLUTIONS - number_adaptation_max
            )
        return solutions

    def get_user_general_solution_nodes(self, user_liberal, user_conservative):
        """Returns a list of general solutions and some information about those general
        solutions. The solutions are ordered from highest to lowest CO2 Equivalent
        Reduced / Sequestered (2020–2050) in Gigatons from Project Drawdown scenario 2.
        """
        general_solutions = self.G.nodes["increase in greenhouse effect"][
            "mitigation solutions"
        ]

        general_solutions_details = []

        for solution in general_solutions:

            current_solution = self.G.nodes[solution]
            self.NX_UTILS.set_current_node(current_solution)
            self.MYTH_PROCESS.set_current_node(current_solution)

            # Filter out nodes that oppose radical political users
            solution_liberal = current_solution["data_properties"]["liberal"]
            solution_conservative = current_solution["data_properties"]["conservative"]

            if user_liberal and solution_liberal == -1:
                continue

            if user_conservative and solution_conservative == -1:
                continue

            d = {
                "iri": self.NX_UTILS.get_node_id(),
                "solutionTitle": current_solution["label"],
                "solutionType": "mitigation",
                "shortDescription": self.NX_UTILS.get_short_description(),
                "longDescription": self.NX_UTILS.get_description(),
                "imageUrl": self.NX_UTILS.get_image_url_or_none(),
                "solutionSpecificMythIRIs": self.MYTH_PROCESS.get_solution_specific_myths(),
                "solutionSources": self.NX_UTILS.get_solution_sources(),
                "solutionCo2EqReduced": self.NX_UTILS.get_co2_eq_reduced(),
            }

            if d not in general_solutions_details:
                general_solutions_details.append(d)

        return general_solutions_details

    def get_user_actions(self, effect_name):
        """
        Takes the name of a climate effect and returns a list of actions associated
        with that node.
        """
        solution_names = self.G.nodes[effect_name]["adaptation solutions"]
        adaptation_solutions = []
        mitigation_solutions = []

        for solution in solution_names:
            try:
                current_solution = self.G.nodes[solution]
                self.NX_UTILS.set_current_node(current_solution)
                self.MYTH_PROCESS.set_current_node(current_solution)

                s_dict = {
                    "iri": self.NX_UTILS.get_node_id(),
                    "solutionTitle": current_solution["label"],
                    "solutionType": "adaptation",
                    "shortDescription": self.NX_UTILS.get_short_description(),
                    "longDescription": self.NX_UTILS.get_description(),
                    "imageUrl": self.NX_UTILS.get_image_url_or_none(),
                    "solutionSpecificMythIRIs": self.MYTH_PROCESS.get_solution_specific_myths(),
                    "solutionSources": self.NX_UTILS.get_solution_sources(),
                }

                # only include direct adaptation solutions (ignore adaptation solutions that are upstream for now)
                for neighbor in self.G.neighbors(effect_name):
                    if (
                        self.G[effect_name][neighbor]["type"]
                        == "is_inhibited_or_prevented_or_blocked_or_slowed_by"
                        and neighbor == solution
                    ):
                        adaptation_solutions.append(s_dict)

            except:
                raise CustomError(
                    message="An error occurred while processing the adaptation solutions for a climate effect in the user's personalized feed."
                )

        solution_names = self.G.nodes["increase in greenhouse effect"][
            "mitigation solutions"
        ]
        for solution in solution_names:
            try:
                self.NX_UTILS.set_current_node(self.G.nodes[solution])
                self.MYTH_PROCESS.set_current_node(self.G.nodes[solution])
                s_dict = {
                    "iri": self.NX_UTILS.get_node_id(),
                    "solutionTitle": self.G.nodes[solution]["label"],
                    "solutionType": "mitigation",
                    "shortDescription": self.NX_UTILS.get_short_description(),
                    "longDescription": self.NX_UTILS.get_description(),
                    "imageUrl": self.NX_UTILS.get_image_url_or_none(),
                    "solutionSpecificMythIRIs": self.MYTH_PROCESS.get_solution_specific_myths(),
                    "solutionSources": self.NX_UTILS.get_solution_sources(),
                }
                mitigation_solutions.append(s_dict)
            except:
                raise CustomError(
                    message="An error occurred while processing the mitigation solutions for a climate effect in the user's personalized feed."
                )

        solutions = self.solution_randomizer(
            adaptation_solutions,
            mitigation_solutions,
        )
        return solutions


"""
Structural component
--------------------

This is an abstract base class for structural components.

"""

import numpy as np
from .component import Component

class ComponentStructural(Component):
    """Abstract base class for structural components.
    """

    @property
    def velocity(self):
        """Velocity of the component :math:`v_{component}`.

        :rtype: :class:`numpy.ndarray`

        This is the sum of all subsystems velocities.
        """
        return sum(subsystem.velocity for subsystem in self.linked_subsystems)

    @property
    def velocity_level(self):
        """
        Velocity level :math:`L_v`.

        :rtype: :class:`numpy.ndarray`

        The structural velocity level is calculated as

        .. math:: L_v = 20 \\log_{10}{\\left( \\frac{v}{v_0} \\right) }

        .. seealso:: :attr:`seapy.system.System.reference_velocity`

        """
        return 20.0 * np.log10(self.velocity / self.system.reference_velocity)


import time
import typing
import warnings
import jax
import jax._src.util as util
import numpy as np
import yaml
import wandb
from jax import lax, numpy as jnp
from jax.experimental import PartitionSpec
from jax.experimental import pjit
from jax.experimental.maps import mesh

from src.context import Context, WhileTrainContext
from src.data import text_dataset
from src.model import compute, body_ctx
from src.optimizer import get_current_lr, update
from src.utils.wandb import WandbLog

def train_step(while_ctx_dict: typing.Dict[str, typing.Any], _unused: None
               ) -> typing.Tuple[typing.Dict[str, typing.Any], None]:
    wctx = WhileTrainContext(while_ctx_dict)
    grad_fn = jax.value_and_grad(compute, 0, True)
    (top_loss, loss), grads = grad_fn(wctx.ctx.parameters,
                                      wctx.data[wctx.current_step % wctx.ctx.training.device_steps])
    update(wctx.ctx, grads, wctx.current_step)
    wctx.loss += loss
    wctx.top_loss += top_loss
    wctx.current_step += 1
    return wctx.serialize(), None

def jitless_step(while_ctx_dict: typing.Dict[str, typing.Any]) -> typing.Dict[str, typing.Any]:
    training = WhileTrainContext(while_ctx_dict).ctx.training
    return lax.scan(train_step, while_ctx_dict, None, training.device_steps, unroll=training.device_unroll)[0]

def sharding(ctx: Context, dims: typing.List[str]):
    out = []
    for d in dims:
        if d == ctx.dims.batch:
            out.append("data_parallel")
        if d == ctx.dims.heads:
            out.append("model_parallel")
        else:
            out.append(None)
    return PartitionSpec(*out)

def timeit(text: str, fn, *args, pad=50):
    start_time = time.time()
    print(f'{text}..', end='', flush=True)
    out = fn(*args)
    print(f"{' ' * (pad - len(text))}Took:{time.time() - start_time:9.2f}s", flush=True)
    return out

def train_loop(wctx: WhileTrainContext, step: typing.Callable):
    ctx = [wctx]

    def _fn(dat: jnp.ndarray) -> WhileTrainContext:
        w = ctx[0](dat)
        w.loss = jnp.zeros_like(w.loss)
        w.top_loss = jnp.zeros_like(w.loss)
        w.current_step = w.current_step + 1
        ctx[0] = WhileTrainContext(step(w.serialize()))
        return ctx[0]

    return _fn

def main():
    warnings.filterwarnings("ignore", message=".*is an experimental feature and probably has bugs!.*")
    # jax.config.update("jax_disable_jit", True)
    wctx = WhileTrainContext()
    ctx = wctx.ctx
    print(yaml.dump(ctx.config(), indent=4))
    ctx.is_initializing = True
    if ctx.wandb.use_wandb:
        run = wandb.init(project=ctx.wandb.project, entity=ctx.wandb.entity,
               config=ctx.config())
        wblog = WandbLog(run)
    total_steps = ctx.training.steps * ctx.training.device_steps
    data = timeit("Initializing dataset", text_dataset, ctx)
    inp = timeit("Enqueueing first batch", next, data)[0, 0]
    timeit("Acquiring forward parameters", body_ctx, ctx, inp)
    parameter_count = sum(util.prod(param.shape) for name, param in ctx.parameters.items())
    timeit("Acquiring optimizer parameters", update, ctx,
           {name: jnp.zeros_like(param) for name, param in ctx.parameters.items()},
           jnp.ones([], dtype=ctx.model.dtype))
    buffer_count = sum(util.prod(param.shape) for name, param in ctx.parameters.items()) - parameter_count

    partition = {'parameters': {name: sharding(ctx, dims) for name, dims in ctx.parameter_dims.items()},
                 'data': PartitionSpec(None, None, "data_parallel", None),
                 'current_step': None, 'loss': None, 'top_loss': None}
    step = train_loop(wctx, timeit("JITing model", pjit.pjit, jitless_step, (partition,), partition))

    mesh_devices = np.array(jax.devices()).reshape(ctx.training.data_parallel, ctx.training.model_parallel)
    global_start = time.time()
    with mesh(mesh_devices, ('data_parallel', 'model_parallel')):
        timeit("Compiling model and performing first step", step, next(data))
        timeit("Running second step", step, next(data))
        print(f"\n\nParameters: {parameter_count:,}\nBuffers:    {buffer_count:,}\n\n")

        start_time = time.time()
        for idx, dat in enumerate(data):
            wctx = step(dat)
            if idx % ctx.training.print_interval == 0:
                millions_processed = ctx.training.device_steps * ctx.dims.sizes.sequence * ctx.dims.sizes.batch
                print(f'[{idx * ctx.training.device_steps:{len(str(total_steps))}d}/{total_steps}] '
                      f'Loss: {wctx.loss / ctx.training.device_steps:6.3f} - '
                      f'TopLoss: {wctx.top_loss / ctx.training.device_steps:8.3f} | '
                      f'LearningRate: {float(get_current_lr(ctx, wctx.current_step)):.5f} | '
                      f'StepTime: {time.time() - start_time:10.6f}s - '
                      f'Rate: {millions_processed * (idx + 1) / (time.time() - global_start):9,.1f} Tokens/s')
            if ctx.wandb.use_wandb and idx % ctx.wandb.log_frequency == 0:
                wblog(wctx, get_current_lr(wctx.ctx, wctx.current_step))
            if ctx.training.trace.do_trace:
                if idx == ctx.training.trace.start_step:
                    jax.profiler.start_trace(ctx.training.trace.output_path)
                if idx == ctx.training.trace.stop_step:
                    jax.profiler.stop_trace()

if __name__ == '__main__':
    main()


from dataset import TorchDataset, TorchDataLoader
from torch.utils.tensorboard import SummaryWriter
import torch
from contrastive_models import ModelWrapper
from losses import bimodal_loss, trimodal_loss
import os
import numpy as np
import shutil
import matplotlib.pyplot as plt

def visualize_data(train_dict, val_dict, writer, epoch):
    fig = plt.figure(figsize=(20, 20))
    gs = fig.add_gridspec(2, 1)
    ax_loss = fig.add_subplot(gs[0])
    ax_acc = fig.add_subplot(gs[1])
    xs = np.arange(epoch + 1)
    for mode, (data_dict, v_tag) in {"Train": (train_dict, '-'), "Val": (val_dict, 'o')}.items():
        losses = np.array(data_dict["losses"])
        accuracies = np.array(data_dict["accuracies"])
        for elt, (c, compare_name) in enumerate([("r", "rgb_audio"), ("g", "rgb_action"),
                                                 ("b", "action_audio"), ("k", "total")]):
            ax_loss.plot(xs, losses[:, elt], f'{c}{v_tag}', label=f"{mode} Loss/{compare_name}")
            if elt != 3:
                ax_acc.plot(xs, accuracies[:, elt], f'{c}{v_tag}', label=f"{mode} Acc/{compare_name}")
    ax_acc.legend()
    ax_loss.legend()
    writer.add_figure("Training Summary", fig, global_step=epoch, close=True)

def train_representations(config):
    """
    Create torch dataset/dataloader objects. Initialize models, optimizer, losses, labels.
    """
    # writer_comment = f"-b:{config['batch_size_train']}-adam:{config['adam']}lr:{config['lr']}-g:{config['gamma']}-m:{config['milestones']}-s:{config['max_samples']}-arch:{config['architecture']}-backbone:{config['cnn_backbone']}-rmode:{config['rgb_mode']}-roffs:{config['rgb_offset']}-tag:{config['tag']}"
    config["num_img_traj"] = 13 - config['rgb_offset']
    writer = SummaryWriter(comment="")
    dataloader = TorchDataLoader(config)
    full_dataloader, train_dataloader, val_dataloader = dataloader.dataloaders
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model_wrapper = ModelWrapper(config)
    model_wrapper.to(device)
    ce_loss = torch.nn.CrossEntropyLoss(), torch.nn.CrossEntropyLoss(), torch.nn.CrossEntropyLoss()
    mse_loss = torch.nn.MSELoss(reduction='none')
    labels_train = torch.arange(config["batch_size_train"]).to(device)
    labels_valid = torch.arange(config["batch_size_valid"]).to(device)
    """
    Run training and validation iterations.
    """
    train_data = {"losses": [], "accuracies": [], "batch_size": config["batch_size_train"]}
    val_data = {"losses": [], "accuracies": [], "batch_size": config["batch_size_valid"]}
    for epoch in range(config["epochs"]):
        train_data["batches"], val_data["batches"] = 0, 0
        train_data["epoch_num_cor"], val_data["epoch_num_cor"] = [0, 0, 0], [0, 0, 0]
        train_data["epoch_losses"], val_data["epoch_losses"] = ([], [], [], []), ([], [], [], [])

        for key in model_wrapper.models.keys():
            [model.train() for model in model_wrapper.models[key].values()]

        for sample_elt, sample in enumerate(train_dataloader):
            if len(sample["rgb"]) < config["batch_size_train"]:
                continue
            if config["architecture"] == "trimodal":
                batch_loss = trimodal_loss(model_wrapper.models, sample, config["batch_size_train"], device, ce_loss,
                                           labels_train, train_data, config)

            elif config["architecture"] == "bimodal":
                batch_loss = bimodal_loss(model_wrapper.models, sample, config["batch_size_train"], device, ce_loss,
                                          labels_train, train_data, config)
            batch_loss.backward()
            for key, optimizer in model_wrapper.optimizers.items():
                optimizer.step()
                optimizer.zero_grad()

        with torch.no_grad():
            for key in model_wrapper.models.keys():
                [model.eval() for model in model_wrapper.models[key].values()]

            for sample_elt, sample in enumerate(val_dataloader):
                if len(sample["rgb"]) < config["batch_size_valid"]:
                    continue

                if config["architecture"] == "trimodal":
                    batch_loss = trimodal_loss(model_wrapper.models, sample, config["batch_size_valid"], device,
                                               ce_loss,
                                               labels_valid, val_data, config)
                elif config["architecture"] == "bimodal":
                    batch_loss = bimodal_loss(model_wrapper.models, sample, config["batch_size_valid"], device, ce_loss,
                                              labels_valid, val_data, config)

        for mode, data_dict in {"Train": train_data, "Val": val_data}.items():
            data_dict["acc_epoch"] = np.array(data_dict["epoch_num_cor"]) / (
                    data_dict["batch_size"] * data_dict["batches"])
            data_dict["avg_loss_epoch"] = np.mean(data_dict["epoch_losses"], axis=1)
            data_dict["losses"].append(data_dict["avg_loss_epoch"])
            data_dict["accuracies"].append(data_dict["acc_epoch"])
            writer.add_scalar(f'Loss/{mode} rgb audio', data_dict["avg_loss_epoch"][0], epoch)
            writer.add_scalar(f'Loss/{mode} rgb actions', data_dict["avg_loss_epoch"][1], epoch)
            writer.add_scalar(f'Loss/{mode} audio actions', data_dict["avg_loss_epoch"][2], epoch)
            writer.add_scalar(f'Loss/{mode} total', data_dict["avg_loss_epoch"][3], epoch)
            writer.add_scalar(f'Accuracy/{mode} rgb audio', data_dict["acc_epoch"][0], epoch)
            writer.add_scalar(f'Accuracy/{mode} rgb actions', data_dict["acc_epoch"][1], epoch)
            writer.add_scalar(f'Accuracy/{mode} audio actions', data_dict["acc_epoch"][2], epoch)
        visualize_data(train_data, val_data, writer, epoch)
        print(train_data["losses"][-1], val_data["losses"][-1], train_data["acc_epoch"], val_data["acc_epoch"])

        if config["save"]:
            model_save_dir = f"{config['architecture']}_models"
            if os.path.exists(model_save_dir):
                shutil.rmtree(model_save_dir)
            os.mkdir(model_save_dir)
            torch.save(model_wrapper.state_dict(), f"{model_save_dir}/wrapper_model.pt")

        [scheduler.step() for scheduler in model_wrapper.schedulers]

if __name__ == "__main__":
        config = {
                "dataset_path": "hw_dataset_craigie", "train_ratio": .8, "max_samples": "max",
                "batch_size_train": 64, "batch_size_valid": 64, "rgb_offset": 4,
                "save": True,
                "cnn_backbone": "resnet", "pretrained": False, "trainable": True,
                "architecture": "trimodal", "adam": True, "epochs": 10,
        }
        train_representations(config)


import PIL
from PIL import Image
import numpy as np

def asampl_apply(img):
    img2 = Image.open('data/sour.png')
    img2 = np.transpose(np.array(img2)[::, ::, :3], axes=[1,0,2])

    img[:img2.shape[0], :img2.shape[1], ::] += img2 // 2
    return img
    # return img / 2 + img2 / 2


import tensorflow as tf
from PIL import Image
import tensorflow as tf
import numpy as np
import cv2
import matplotlib.pyplot as plt

class OD:
    def __init__(self):
        config = tf.ConfigProto(allow_soft_placement=True)
        config.gpu_options.allow_growth = True
        config.gpu_options.per_process_gpu_memory_fraction = 1

        detection_graph = tf.Graph()
        with detection_graph.device('/gpu:' + str(0)):
            with detection_graph.as_default():
                od_graph_def = tf.GraphDef()
                with tf.gfile.GFile('frozen_inference_graph.pb', 'rb') as fid:
                    serialized_graph = fid.read()
                    od_graph_def.ParseFromString(serialized_graph)
                    tf.import_graph_def(od_graph_def, name='')

                self.sess = tf.Session(graph=detection_graph, config=config)
                # Fetching tensors from the graph
                self.image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
                self.detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
                self.detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
                self.detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
                self.num_detections = detection_graph.get_tensor_by_name('num_detections:0')

    def infer(self, image):
        image_np_expanded = np.reshape(image, (1, 540, 960, 3))
        (boxes, scores, classes, num) = self.sess.run(
            [self.detection_boxes, self.detection_scores, self.detection_classes, self.num_detections],
            feed_dict={self.image_tensor: image_np_expanded})
        scores = scores[0]
        indexs = np.where(scores >= 0.90)
        for index in indexs[0]:
            x1 = int(540 * boxes[0][index][0])
            y1 = int(960 * boxes[0][index][1])
            x2 = int(540 * boxes[0][index][2])
            y2 = int(960 * boxes[0][index][3])
            cv2.rectangle(image, (y1, x1), (y2, x2), (0, 255, 0), 2)
        return image, len(indexs[0])


#!/usr/bin/env python
# encoding: utf-8
# author: 04

import operator
import itertools
from decimal import Decimal as D # noqa
from functools import reduce
from queue import PriorityQueue as PQ # noqa

import yaml
import click
from scipy.special import comb
from tqdm import tqdm

from buildings import BuildingMatrix, GlobalBuffer
from consts import BUILDING_INFO, BufferConsts as Bc
from errors import MatrixCategoryFull, MatrixFull

CUSTOM_FILE_NAME = 'jiaguomeng.yml'

class CalcJiaGuoMeng(object):
    """ 主函数，收集用户信息并完成计算
    """

    def __init__(self, online_mod, conf, only):
        self.building_matrix = BuildingMatrix(BUILDING_INFO)
        self.custom_config = self.read_custom_config(conf)
        self.online_mod = online_mod
        self.only_one_building = only
        self.fill_global_buffer()

    def read_custom_config(self, conf_file):
        config = yaml.load(open(conf_file), Loader=yaml.SafeLoader)
        for config_name in ['1星', '2星', '3星', '4星', '5星']:
            bnames = config[config_name] or ''
            star = int(config_name[:1])
            for name in bnames.split():
                self.building_matrix.buildings[name].set_star(star)
        # 星级全部确定后，可确定所有加成具体数值，即可对buff排序
        self.building_matrix.sort_buffer()
        effect_trans = {
            '政策': 'policy',
            '照片': 'photos',
            '城市任务': 'quests'
        }
        self.global_effects = {
            effect_name: self._read_custom_buffers(effect_name, config[config_name])
            for config_name, effect_name in effect_trans.items()
        }
        self.global_effects['quests'].extend(
            self._read_custom_binds('quests', config['城市任务建筑加成'])
        )

    def _read_custom_buffers(self, global_type, buffer_conf):
        buffer_trans = {
            '在线': Bc.ONL,
            '离线': Bc.OFL,
            '住宅': Bc.RES,
            '商业': Bc.COM,
            '工业': Bc.IND,
        }
        return [
            GlobalBuffer(global_type, buffer_type, D(buffer_conf.get(conf_key)))
            for conf_key, buffer_type in buffer_trans.items()
            if buffer_conf.get(conf_key) is not 0
        ]

    def _read_custom_binds(self, global_type, bind_config):
        return [
            GlobalBuffer(global_type, Bc.SGL, D(coefficient), building_name)
            for building_name, coefficient in bind_config.items()
            if coefficient is not 0
        ]

    def fill_global_buffer(self):
        for building in self.building_matrix.buildings.values():
            match_effects = {
                effects_name: list(filter(
                    lambda buf: buf.fit_income(building, online=self.online_mod), buffer_list
                ))
                for effects_name, buffer_list in self.global_effects.items()
            }
            global_coefficient = reduce(operator.mul, [
                D(1) + sum([buffer.coefficient for buffer in buffer_list])
                for buffer_list in match_effects.values()
            ])
            building.global_coefficient = global_coefficient

class ExplosionMixin(object):
    """ 暴力破解，高消耗
    仅用作验证，不做实际运行使用，未做数据结构简化的优化
    """

    def print_plan(self, total_score, main_bd, plan):
        print('总加成: {}'.format(total_score))
        print('主建筑: {}'.format(main_bd.name))
        print('建筑列表:')
        if isinstance(plan, dict):
            plan = [plan[line] for line in [Bc.RES, Bc.COM, Bc.IND]]
        for index, cn in enumerate(['住宅', '商业', '工业']):
            print('{idt}{cn}:{bds}'.format(
                idt=' ' * 8,
                cn=cn,
                bds=' '.join(bd.name for bd in plan[index])
            ))

    def explosion(self):
        """ 网上的暴力破解方法
        result 最大范围限定为3，随时取出最后一个，即保证前两个是最优化和次优化配置
        """
        _m = self.building_matrix
        res = _m.indexes[Bc.RES]
        com = _m.indexes[Bc.COM]
        ind = _m.indexes[Bc.IND]
        search_space = itertools.product(
            itertools.combinations(res, 3),
            itertools.combinations(com, 3),
            itertools.combinations(ind, 3)
        )
        search_space_size = comb(len(ind), 3) * comb(
            len(com), 3) * comb(len(res), 3)
        print('Total iterations:', search_space_size)
        results = PQ()
        for plan in tqdm(
                search_space,
                total=search_space_size,
                bar_format='{percentage:3.0f}%, {elapsed}<{remaining}|{bar}|{n_fmt}/{total_fmt}, {rate_fmt}{postfix}',
                ncols=80):
            prod = self.explosion_calc(plan)
            results.put((-1 * prod[0], prod[1], plan))
        first_good = results.get()
        second_good = results.get()
        print('The First Good Plan')
        self.print_plan(*first_good)
        print('The Second Good Plan')
        self.print_plan(*second_good)

    def explosion_calc(self, plan):
        plan_buildings = [bd for cat in plan for bd in cat]
        for bd in plan_buildings:
            active_buffers = [
                buffer for buffer in bd.buffed_by
                if buffer.buffer_from in plan_buildings and
                buffer.fit_income(bd)
            ]
            bd.result = (
                D(1) + sum(buf.coefficient for buf in active_buffers)
            ) * bd.global_coefficient
        priority_order = sorted(plan_buildings, key=lambda x: x.result)
        return sum(bd.result for bd in plan_buildings), priority_order[-1]

class PickUpMixin(object):
    """ 按人为思路筛选
    基于最低成本考虑，主建筑总是唯一一个，允许有次升建筑。
    """
    debug = {}

    def run(self):
        """ 计算并寻找最优解
        计算方式：
        1. 首先确定每个建筑全局加成
        2. 计算每个建筑最高加成建筑的组合,及加成总倍率,确定主力建筑
        3. 合并次要方案，并计算总体建筑价值，基于建筑价值排序去除低价值建筑。
        总体建筑价值 = 建筑直接收益系数 + 建筑加成间接收益系数
        4. 输出方案及升级价值排序。
        升级价值 == 建筑直接收益系数
        """
        building_plans = self.first_building_plans()
        self.helper_buildings = [info['bd'] for info in building_plans[30:]]
        main_plan = building_plans[0]
        if self.only_one_building:
            # 一般主力建筑等级高出其他建筑50~100级，高出往期主力建筑20~50级，主力建筑系数调整5倍
            main_plan['bd'].global_coefficient *= 5
        # 执行三次优化, 逐步扩大搜索范围，下探到较差的组合中确认是否有互补情况
        total_income, counted_detail = self.count_total_income(main_plan['plan'])
        consider_plan = confirmed_plan = {
            'plan': main_plan['plan'],
            'total_income': total_income,
            'detail': counted_detail
        }
        seek_ranges = [4, 7, 10]
        for seek_range in seek_ranges:
            value_plans = building_plans[1:seek_range]
            for pick_plan in value_plans:
                merged_plan = self.merge_plans(confirmed_plan['plan'], pick_plan['plan'])
                merge_income, merge_detail = self.count_total_income(merged_plan)
                if merge_income > consider_plan['total_income']:
                    consider_plan = {
                        'plan': merged_plan,
                        'total_income': merge_income,
                        'detail': merge_detail
                    }
            confirmed_plan = consider_plan
        confirmed_plan['sorted_detail'] = sorted([
            (bd, value['direct_income'], value['indirect_income'])
            for bd, value in confirmed_plan['detail'].items()
        ], key=lambda x: x[1] + x[2], reverse=True)
        self.print_plan(confirmed_plan)

    def first_building_plans(self):
        calc_completed = []
        ordered_building_list = sorted(
            self.building_matrix.buildings.values(),
            key=lambda bd: bd.global_coefficient * bd.self_effect,
            reverse=True
        )
        for bd in ordered_building_list:
            plan = self.building_matrix.put(bd)
            building_effect = D(1)
            ordered_buffer_buildings = sorted(bd.buffed_by, key=lambda x: x.coefficient, reverse=True)
            for buffer in ordered_buffer_buildings:
                if not buffer.fit_income(bd, online=self.online_mod):
                    continue
                try:
                    plan = self.building_matrix.put(buffer.buffer_from, plan)
                except MatrixCategoryFull:
                    continue
                except MatrixFull:
                    break
                building_effect += buffer.coefficient
            calc_completed.append({
                'bd': bd,
                'max_bd_effect': building_effect,
                'income_coefficient': bd.global_coefficient * bd.self_effect * building_effect,
                'plan': plan,
            })
        calc_completed.sort(key=lambda x: x['income_coefficient'], reverse=True)
        return calc_completed

    def count_total_income(self, plan, explain=False):
        flat_plan = {
            bd: {
                'direct_income': 0,
                'indirect_income': 0,
            } for line in plan.values() for bd in line}
        explain_data = {}
        for bd, info in flat_plan.items():
            if bd in self.helper_buildings:
                # 辅助建筑只计算间接受益
                continue
            bd_buffed = D(1)
            bd_baseline = bd.global_coefficient * bd.self_effect
            for buf in bd.buffed_by:
                if not buf.fit_income(bd, online=self.online_mod):
                    continue
                if buf.buffer_from in flat_plan:
                    bd_buffed += buf.coefficient
                    flat_plan[buf.buffer_from]['indirect_income'] += (
                        bd_baseline * buf.coefficient)
                    if explain:
                        effect_num = bd_baseline * buf.coefficient
                        bd_explain = explain_data.setdefault(bd, {'buffed_from': [], 'buffer_to': []})
                        bd_explain['buffed_from'].append((buf.buffer_from, effect_num))
                        buffer_from_explain = explain_data.setdefault(buf.buffer_from, {'buffed_from': [], 'buffer_to': []})
                        buffer_from_explain['buffer_to'].append((bd, effect_num))
                        self.debug['explain'] = explain_data
            info['direct_income'] = bd_baseline * bd_buffed
        total_income = sum([info['direct_income'] for info in flat_plan.values()])
        return total_income, flat_plan

    def merge_plans(self, main_plan, pick_plan):
        """ 尝试合并两个plan, 合并后剔除最低加成的建筑
        """
        merging_plan = {key: set.union(main_plan[key], pick_plan[key]) for key in main_plan}
        total_income, detail = self.count_total_income(merging_plan, explain=True)
        for line_name, bd_set in merging_plan.items():
            if len(bd_set) > 3:
                last_bds = sorted(
                    list(bd_set),
                    key=lambda x: detail[x]['direct_income'] + detail[x]['indirect_income'],
                    reverse=True
                )[3:]
                merging_plan[line_name] = bd_set.difference(set(last_bds))
        return merging_plan

    def print_plan(self, confirmed_plan):
        """
        """
        print("建筑方案：")
        bd_types = [('住宅', Bc.RES), ('商业', Bc.COM), ('工业', Bc.IND)]
        for cn, bd_type in bd_types:
            print(f'''{cn}: {''.join([
                '{:<{len}}'.format(bd.name, len=10-len(bd.name)) for bd in confirmed_plan['plan'][bd_type]
            ])}''')
        print('=' * 80)
        titles = ['建筑名称', '直接收益', '间接收益']
        cell_width = 10
        print('建筑价值')
        detail = confirmed_plan['sorted_detail']
        for key in range(3):
            if titles[key] == '建筑名称':
                line = ''.join(['{value:<{len}}'.format(value=value[key].name, len=cell_width - len(value[key].name))
                                for value in detail])
            else:
                line = ''.join(['{value:<{len}.2f}'.format(value=value[key], len=cell_width) for value in detail])
            print(f'''{titles[key]:<4}: {line}''')
        print('=' * 80)
        print('升级优先度')
        print(''.join(
            '{name:<{len}}'.format(
                name=bd.name, len=cell_width - len(bd.name)
            ) for bd, direct, indirect in sorted(detail, key=lambda x:x[1], reverse=True)
            if direct >= indirect
        ))
        print('=' * 80)
        print('总系数')
        print(sum([direct for bd, direct, indirect in detail]))

class CalcByExplosion(CalcJiaGuoMeng, ExplosionMixin):
    pass

class CalcByPick(CalcJiaGuoMeng, PickUpMixin):
    pass

@click.command()
@click.option('-b', '--bomm', is_flag=True, help='Use Explosion Mod.')
@click.option('-f', '--offline', is_flag=True, help='Offline Mod.')
@click.option('-o', '--only', is_flag=True, help='Only One Building is very important!')
@click.option('-c', '--config', default=CUSTOM_FILE_NAME, help='Set conf file path.')
def main(bomm, offline, config, only):
    args = {
        'online_mod': not offline,
        'conf': config,
        'only': only,
    }
    if bomm:
        calculater = CalcByExplosion(**args)
        calculater.explosion()
    else:
        calculater = CalcByPick(**args)
        calculater.run()

if __name__ == '__main__':
    main()


import os
import numpy as np
import tensorrt as trt
import trt_utils
import time
import alphabets_english as alp
import dataset
from PIL import Image
TRT_LOGGER = trt.Logger()
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

img_path = './demo/img_14.jpg'
engine_path = './new_crnn_english.trt'
batch = 10
width = 1024
alphabets = alp.alphabet

batch_images = np.zeros(shape=(32, 1, 32, width), dtype=np.float32)
image = Image.open(img_path).convert('L')
image = image.resize([int(width/1.5),32])
image = np.array(image, np.float32)
image = image / 255
for i in range(batch):
    batch_images[i, 0, :, :int(width/1.5)] = image

def decode(resd):
    _, w = resd.shape
    texts = []
    for i in range(batch):
        l = resd[i]
        char_list=[]
        for i in range(w):
            if l[i] != 0 and (not (i > 0 and l[i - 1] == l[i])):
                char_list.append(alphabets[l[i] - 1])
        texts.append(''.join(char_list))
    return texts

with trt_utils.get_engine(
        engine_path, TRT_LOGGER) as engine, engine.create_execution_context() as context:

    inputs, outputs, bindings, stream = trt_utils.allocate_buffers(engine)
    # batch_images = np.random.random_sample([32, 1, 32, 1024])
    for i in range(10):
        start = time.time()
        inputs[0].host = np.array(batch_images, dtype=np.float32,
                                  order='C')
        trt_outputs = trt_utils.do_inference_v2(context, bindings=bindings,
                                             inputs=inputs, outputs=outputs,
                                             stream=stream)
        trt_outputs = trt_outputs[0].reshape([int(width/4+1), 32, len(alphabets)+1])
        trt_outputs = np.argmax(trt_outputs, axis=2)
        # print(len(batch_outputs))
        trt_outputs = np.transpose(trt_outputs,(1,0))[:batch, :]
        print(decode(trt_outputs))
        end = time.time()
        print(end - start)


# *-* coding: utf-8 *-*
"""Manage measurement configurations and corresponding measurements.

Geometric Factors
-----------------

.. plot:: pyplots/plot_K_vs_dipol_sep.py

Literature
----------

* Noel, Mark and Biwen Xu; Archaeological investigation by electrical
    resistivity tomography: a preliminary study. Geophys J Int 1991; 107 (1):
    95-102. doi: 10.1111/j.1365-246X.1991.tb01159.x
* Stummer, Peter, Hansruedi Maurer, and Alan G. Green. “Experimental Design:
  Electrical Resistivity Data Sets That Provide Optimum Subsurface
  Information.” Geophysics 69, no. 1 (January 1, 2004): 120–120.
  doi:10.1190/1.1649381.
* Roy, A., and A. Apparao. “DEPTH OF INVESTIGATION IN DIRECT CURRENT METHODS.”
  GEOPHYSICS 36, no. 5 (January 1, 1971): 943–59. doi:10.1190/1.1440226.

"""
import numpy as np
import pandas as pd

import crtomo.mpl

import reda.configs.configManager as reda_config_mgr
plt, mpl = crtomo.mpl.setup()

class ConfigManager(reda_config_mgr.ConfigManager):
    def __init__(self, **kwargs):
        """

        nr_of_electrodes : int
            Number of electrodes
        """
        # assert 'nr_of_electrodes' in kwargs
        super().__init__(**kwargs)
        # store measurements in a list of size N arrays
        self.measurements = {}
        # global counter for measurements
        self.meas_counter = - 1

    def clear_measurements(self):
        """Remove all measurements from self.measurements. Reset the
        measurement counter. All ID are invalidated.
        """
        keys = list(self.measurements.keys())
        for key in keys:
            del(self.measurements[key])
        self.meas_counter = -1

    def delete_measurements(self, mid):
        del(self.measurements[mid])

    def add_measurements(self, measurements):
        """Add new measurements to this instance

        Parameters
        ----------
        measurements: numpy.ndarray
            one or more measurement sets. It must either be 1D or 2D, with the
            first dimension the number of measurement sets (K), and the second
            the number of measurements (N): K x N

        Returns
        -------
        mid: int
            measurement ID used to extract the measurements later on

        Examples
        --------
        >>> import numpy as np
            import crtomo.configManager as CRconfig
            config = CRconfig.ConfigManager(nr_of_electrodes=10)
            config.gen_dipole_dipole(skipc=0)
            # generate some random noise
            random_measurements = np.random.random(config.nr_of_configs)
            mid = config.add_measurements(random_measurements)
            # retrieve using mid
            print(config.measurements[mid])

        """
        subdata = np.atleast_2d(measurements)

        if self.configs is None:
            raise Exception(
                'must read in configuration before measurements can be stored'
            )

        # we try to accommodate transposed input
        if subdata.shape[1] != self.configs.shape[0]:
            if subdata.shape[0] == self.configs.shape[0]:
                subdata = subdata.T
            else:
                raise Exception(
                    'Number of measurements does not match number of configs'
                )

        return_ids = []
        for dataset in subdata:
            cid = self._get_next_index()
            self.measurements[cid] = dataset.copy()
            return_ids.append(cid)

        if len(return_ids) == 1:
            return return_ids[0]
        else:
            return return_ids

    def remove_using_nr_injections(self, min_nr_of_injections):
        """Remove all measurements with a current injection that is not used a
        minimum number of times.

        This is useful to optimize measurement time for multi-channel systems,
        such as the IRIS Instruments Syscal Pro. The device can theoretically
        only facilitate the full number of channels if a corresponding number
        of voltage measurements is requested for a given current injection. As
        such, it can be useful to remove measurements with unique current
        injection dipoles. Note that other factors determine if all channels
        are actually used. Please refer to the device manual for
        device-specific information.

        Parameters
        ----------
        min_nr_of_injections : int
            Minimum number a given current injection should have to keep it

        """
        if min_nr_of_injections <= 0:
            return

        df = pd.DataFrame(self.configs)
        # group over A, B
        g = df.groupby([0, 1])
        df_filtered = g.filter(lambda x: x.shape[0] >= min_nr_of_injections)
        self.configs = df_filtered.values

    def plot_error_pars(self, mid):
        """ ??? DEFUNCT

        """
        R = None
        fig, axes = plt.subplots(1, 2, figsize=(10, 6))

        def plot_error_pars(axes, a, b, R, label=''):
            dR = a * R + b
            dlogR = np.abs(a + b / R)
            ax = axes[0]
            ax.scatter(R, dR / R * 100, label=label)
            ax = axes[1]
            ax.scatter(R, dlogR / np.abs(np.log(R)) * 100, label=label)
            ax.set_xscale('log')
            ax.set_yscale('log')

        for b in np.linspace(R.min(), np.percentile(R, 10), 5):
            plot_error_pars(axes, 0.05, b, R, label='b={0:.4f}'.format(b))
        axes[0].set_xlabel(r'$R [\Omega]$')
        axes[0].set_ylabel(r'$(\Delta R/R) \cdot 100 [\%]$')
        axes[1].set_xlabel(r'$log_{10}(R [\Omega])$')
        axes[1].set_ylabel(
            r'$log_{10}(R) / \Delta log_{10}(R) \cdot 100 [\%]$'
        )
        axes[0].axhline(y=100)
        axes[1].axhline(y=100)
        fig.tight_layout()
        fig.subplots_adjust(bottom=0.2)
        axes[0].legend(
            loc="lower center",
            ncol=4,
            bbox_to_anchor=(0, 0, 1, 1),
            bbox_transform=fig.transFigure
        )

        # fig.savefig('out.png', dpi=300)
        return fig

    def gen_voltage_dipoles_skip(self, cinj, skip=0):
        """For given current injections, generate all possible voltage dipoles
        with skip **skip**.

        Parameters
        ----------
        cinj : :py:class:`numpy.ndarray`
            Nx2 array containing the current injection electrodes a,b
        skip : int, optional
            Skip used for voltage electrodes. Default: 0

        """
        new_configs = []
        for ab in cinj:
            for i in range(1, self.nr_electrodes + 1):
                m = i
                n = i + skip + 1
                if m in ab or n in ab:
                    continue
                if n > self.nr_electrodes:
                    continue
                new_configs.append((ab[0], ab[1], m, n))
        configs = np.array(new_configs)
        self.add_to_configs(configs)
        return configs

    def load_crmod_data(self, data_source):
        """Load CRMod measurement data, either from file or directly from a
        numpy array.

        Parameters
        ----------
        data_source : string|numpy.ndarray
            if this is a string, treat it as an input filename. If it is a Nx5
            or Nx3 numpy array, use this data directly
        Returns
        -------
        cid_mag : int
            Measurement id for magnitude data
        cid_pha : int
            Measurement id for phase data
        """
        if isinstance(data_source, str):
            with open(data_source, 'r') as fid:
                nr_of_configs = int(fid.readline().strip())
                measurements = np.loadtxt(fid)
                if nr_of_configs != measurements.shape[0]:
                    raise Exception(
                        'indicated number of measurements does not equal '
                        'to actual number of measurements'
                    )
        elif isinstance(data_source, pd.DataFrame):
            measurements = data_source[
                ['a', 'b', 'm', 'n', 'r', 'rpha']
            ].values
        else:
            # assume numpy array
            # data already stored in the variable
            measurements = data_source

        assert measurements.shape[1] in (4, 6)
        if measurements.shape[1] == 4:
            # crmod ab-mn scheme
            abmn = self._crmod_to_abmn(measurements[:, 0:2])
            rmag = measurements[:, 2]
            rpha = measurements[:, 3]
        else:
            # abmn in separate columns
            abmn = measurements[:, 0:4]
            rmag = measurements[:, 4]
            rpha = measurements[:, 5]

        if self.configs is None:
            self.configs = abmn
        else:
            # check that configs match
            if not np.all(abmn == self.configs):
                raise Exception(
                    'previously stored configurations do not match new '
                    'configurations'
                )
        # add data
        cid_mag = self.add_measurements(rmag)
        cid_pha = self.add_measurements(rpha)
        return cid_mag, cid_pha

    def load_crmod_volt(self, filename):
        """Load a CRMod measurement file (commonly called volt.dat)

        Parameters
        ----------
        filename : string
            path to input filename

        Returns
        -------
        cid_mag : int
            Measurement id for magnitude data
        cid_pha : int
            Measurement id for phase data
        """
        cid_mag, cid_pha = self.load_crmod_data(filename)
        return cid_mag, cid_pha

    def delete_data_points(self, indices):
        """Delete data points by index (0-indexed), both in configs and
        measurements. Deletions will be done in ALL registered measurements to
        ensure consistency.

        Parameters
        ----------
        indices : int|iterable
            Indices to delete
        """
        print('Deleting configurations:')
        print(self.configs[indices])

        # first the configurations
        self.configs = np.delete(self.configs, indices, axis=0)

        for key in sorted(self.measurements.keys()):
            self.measurements[key] = np.delete(
                self.measurements[key], indices, axis=0
            )


import numpy as np
import argparse
from experiment.uncertainty import uncertainty
from plots.rec_plots import show_uncertainty
from utils.io import load_numpy, save_dataframe_csv, find_best_hyperparameters, load_yaml

def main(args):
    table_path = load_yaml('config/global.yml', key='path')['tables']
    df = find_best_hyperparameters(table_path+args.problem, 'NDCG')
    R_train = load_numpy(path=args.path, name=args.train)
    results = uncertainty(R_train, df, rank=200)

    show_uncertainty(results, x='numRated', y='std', hue='model',
                     folder=args.problem, name=args.name, save=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ParameterTuning")
    parser.add_argument('-n', dest='name', default="uncertainty_analysis")
    parser.add_argument('-d', dest='path', default="datax/")
    parser.add_argument('-t', dest='train', default='Rtrain.npz')
    parser.add_argument('-p', dest='problem', default='movielens1m')
    args = parser.parse_args()

    main(args)


import os

import numpy as np
import cv2

# Read the query image.
folder = 'tattoos'
query = cv2.imread(os.path.join(folder, 'query.png'),
                   cv2.IMREAD_GRAYSCALE)

# create files, images, descriptors globals
files = []
images = []
descriptors = []
for (dirpath, dirnames, filenames) in os.walk(folder):
    files.extend(filenames)
    for f in files:
        if f.endswith('npy') and f != 'query.npy':
            descriptors.append(f)
print(descriptors)

# Create the SIFT detector.
sift = cv2.xfeatures2d.SIFT_create()

# Perform SIFT feature detection and description on the
# query image.
query_kp, query_ds = sift.detectAndCompute(query, None)

# Define FLANN-based matching parameters.
FLANN_INDEX_KDTREE = 1
index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
search_params = dict(checks=50)

# Create the FLANN matcher.
flann = cv2.FlannBasedMatcher(index_params, search_params)

# Define the minimum number of good matches for a suspect.
MIN_NUM_GOOD_MATCHES = 10

greatest_num_good_matches = 0
prime_suspect = None

print('>> Initiating picture scan...')
for d in descriptors:
    print('--------- analyzing %s for matches ------------' % d)
    matches = flann.knnMatch(
        query_ds, np.load(os.path.join(folder, d)), k=2)
    good_matches = []
    for m, n in matches:
        if m.distance < 0.7 * n.distance:
            good_matches.append(m)
    num_good_matches = len(good_matches)
    name = d.replace('.npy', '').upper()
    if num_good_matches >= MIN_NUM_GOOD_MATCHES:
        print('%s is a suspect! (%d matches)' % \
            (name, num_good_matches))
        if num_good_matches > greatest_num_good_matches:
            greatest_num_good_matches = num_good_matches
            prime_suspect = name
    else:
        print('%s is NOT a suspect. (%d matches)' % \
            (name, num_good_matches))

if prime_suspect is not None:
    print('Prime suspect is %s.' % prime_suspect)
else:
    print('There is no suspect.')