import numpy as np
import math
import time
import cv2
import os
import sys
import collections
from datetime import datetime
import random
from utils import *

class HeartrateEstimator(object):
    def __init__(self, pixel_threshold=1.0):

        self.fps = 0
        self.buffer_size = 250
        self.essential_frame_accumulated = 30
        self.data_buffer_a = collections.deque(maxlen=self.buffer_size)
        self.data_buffer_b = collections.deque(maxlen=self.buffer_size)
        self.current_buffer = "a"
        self.times_a = collections.deque(maxlen=self.buffer_size)
        self.times_b = collections.deque(maxlen=self.buffer_size)

        self.freqs = []
        self.fft = []
        self.t0 = time.time()
        self.bpms = circularlist(self.essential_frame_accumulated)
        self.bpm = 0
        self.pixel_threshold = pixel_threshold

        self.face_rect = None

        self.find_faces = True

        self.icon_size = 30
        self.text_icon_size = 60

        self.pre_alpha = 0
        self.pre_reverse = False

        self.heart_img = cv2.imread("./imgs/heart.png", cv2.IMREAD_UNCHANGED)

        self.heart_img = cv2.resize(self.heart_img, (self.icon_size, self.icon_size))

        self.bpm_img = cv2.imread("./imgs/bpm.png", cv2.IMREAD_UNCHANGED)
        self.bpm_img = cv2.resize(
            self.bpm_img, (self.text_icon_size, self.text_icon_size)
        )

        self.ft = cv2.freetype.createFreeType2()
        self.ft.loadFontData(fontFileName="./fonts/HelveticaNeue.ttf", id=0)

    # Draw disconnected rectangle for a bounding box region
    def draw_disconnected_rect(self, img, pt1, pt2, color, thickness):
        width = pt2[0] - pt1[0]
        height = pt2[1] - pt1[1]
        line_width = min(20, width // 4)
        line_height = min(20, height // 4)
        line_length = max(line_width, line_height)
        cv2.line(img, pt1, (pt1[0] + line_length, pt1[1]), color, thickness)
        cv2.line(img, pt1, (pt1[0], pt1[1] + line_length), color, thickness)
        cv2.line(
            img, (pt2[0] - line_length, pt1[1]), (pt2[0], pt1[1]), color, thickness
        )
        cv2.line(
            img, (pt2[0], pt1[1]), (pt2[0], pt1[1] + line_length), color, thickness
        )
        cv2.line(
            img, (pt1[0], pt2[1]), (pt1[0] + line_length, pt2[1]), color, thickness
        )
        cv2.line(
            img, (pt1[0], pt2[1] - line_length), (pt1[0], pt2[1]), color, thickness
        )
        cv2.line(img, pt2, (pt2[0] - line_length, pt2[1]), color, thickness)
        cv2.line(img, (pt2[0], pt2[1] - line_length), pt2, color, thickness)

    def find_faces_toggle(self):
        self.find_faces = not self.find_faces
        return self.find_faces

    # Otain the scaled coordinates of a region with normalized coordinates
    def get_subface_coord(self, fh_x, fh_y, fh_w, fh_h):
        x, y, w, h = self.face_rect
        return [
            int(x + w * fh_x - (w * fh_w / 2.0)),
            int(y + h * fh_y - (h * fh_h / 2.0)),
            int(w * fh_w),
            int(h * fh_h),
        ]

    # Get the mean values for B, G, R channels of an image
    def get_subface_means(self, frame, coord):
        x, y, w, h = coord
        subframe = frame[y : y + h, x : x + w, :]
        v1 = np.mean(subframe[:, :, 0])
        v2 = np.mean(subframe[:, :, 1])
        v3 = np.mean(subframe[:, :, 2])
        return v1, v2, v3

    def select_storing_buffer(self, current_value, threshold):
        selected_buffer = None
        selected_times = None
        if self.current_buffer == "a":
            selected_buffer = self.data_buffer_a
            selected_times = self.times_a
        else:
            selected_buffer = self.data_buffer_b
            selected_times = self.times_b

        if len(selected_buffer) > 0:
            diff = abs(current_value - selected_buffer[-1])
            if diff > threshold:
                print(
                    "Change exceeds pixel threshold, may need to switch to new data buffer"
                )
                if self.current_buffer == "a":
                    if len(self.data_buffer_a) < self.buffer_size:
                        # No need to switch buffer, just need to clear existing data of current buffer
                        self.data_buffer_a = collections.deque(maxlen=self.buffer_size)
                        self.times_a = collections.deque(maxlen=self.buffer_size)
                        selected_buffer = self.data_buffer_a
                        selected_times = self.times_a
                    else:
                        # Switch to another buffer to preserve the current steady signal
                        self.current_buffer = "b"
                        self.data_buffer_b = collections.deque(maxlen=self.buffer_size)
                        self.times_b = collections.deque(maxlen=self.buffer_size)
                        selected_buffer = self.data_buffer_b
                        selected_times = self.times_b

                else:
                    if len(self.data_buffer_b) < self.buffer_size:
                        # No need to switch buffer, just need to clear existing data of current buffer
                        self.data_buffer_b = collections.deque(maxlen=self.buffer_size)
                        self.times_b = collections.deque(maxlen=self.buffer_size)
                        selected_buffer = self.data_buffer_b
                        selected_times = self.times_b
                    else:
                        # Switch to another buffer to preserve the current steady signal
                        self.current_buffer = "a"
                        self.data_buffer_a = collections.deque(maxlen=self.buffer_size)
                        self.times_a = collections.deque(maxlen=self.buffer_size)
                        selected_buffer = self.data_buffer_a
                        selected_times = self.times_a
        return selected_buffer, selected_times

    def update_storing_buffer(self, buffer, times):
        if self.current_buffer == "a":
            self.data_buffer_a = buffer
            self.times_a = times
        else:
            self.data_buffer_b = buffer
            self.times_b = times

    def calculate_bpm(self, buffer, times):
        processed = np.array(buffer)
        bpm = 0
        L = len(buffer)
        if L > self.essential_frame_accumulated:
            fps = float(L) / (times[-1] - times[0])
            even_times = np.linspace(times[0], times[-1], L)
            interpolated = np.interp(even_times, times, processed)
            interpolated = np.hamming(L) * interpolated
            interpolated = interpolated - np.mean(interpolated)
            raw = np.fft.rfft(interpolated)
            phase = np.angle(raw)
            fft = np.abs(raw)
            freqs = float(fps) / L * np.arange(L / 2 + 1)
            freqs = 60.0 * freqs
            idx = np.where((freqs > 50) & (freqs < 240))
            pruned = fft[idx]
            phase = phase[idx]

            pfreq = freqs[idx]
            freqs = pfreq
            fft = pruned
            idx2 = np.argmax(pruned)

            gap = (self.buffer_size - L) / fps
            if gap:
                bpm = 0
            else:
                bpm = freqs[idx2]
        return bpm

    def draw_object_imgs(self, image, object_img, x1, y1, x2, y2, alpha):
        if x1 >= 0 and y1 >= 0 and x2 < image.shape[1] and y2 < image.shape[0]:
            object_alpha = object_img[:, :, 3] / 255.0
            combined_alpha = object_alpha * alpha
            y2 = y2 + (object_img.shape[0] - (y2 - y1))
            image[y1:y2, x1:x2, 0] = (1.0 - combined_alpha) * image[
                y1:y2, x1:x2, 0
            ] + combined_alpha * object_img[:, :, 0]
            image[y1:y2, x1:x2, 1] = (1.0 - combined_alpha) * image[
                y1:y2, x1:x2, 1
            ] + combined_alpha * object_img[:, :, 1]
            image[y1:y2, x1:x2, 2] = (1.0 - combined_alpha) * image[
                y1:y2, x1:x2, 2
            ] + combined_alpha * object_img[:, :, 2]

    def draw_fade_heart(self, frame, x, y, w, h):
        self.draw_object_imgs(
            frame,
            self.heart_img,
            x + w // 2 - self.icon_size // 2,
            y + int(0.2 * h) - self.icon_size // 2,
            x + w // 2 - self.icon_size // 2 + self.icon_size,
            y + int(0.2 * h) - self.icon_size // 2,
            self.pre_alpha,
        )
        if not self.pre_reverse:
            self.pre_alpha = self.pre_alpha + 0.05
            if self.pre_alpha >= 1:
                self.pre_reverse = True
        else:
            self.pre_alpha = self.pre_alpha - 0.05
            if self.pre_alpha <= 0:
                self.pre_reverse = False

    def draw_pump_heart(self, frame, x, y, alpha):
        scale_factor = (((alpha - 0) * 0.5) / 1) + 0.5
        heart_img = cv2.resize(
            self.heart_img,
            (0, 0),
            fx=scale_factor,
            fy=scale_factor,
            interpolation=cv2.INTER_NEAREST,
        )
        x1 = x - self.icon_size - 70 + (self.icon_size // 2 - heart_img.shape[1] // 2)
        y1 = y + 10 + (self.icon_size // 2 - heart_img.shape[0] // 2)
        x2 = x1 + heart_img.shape[1]
        y2 = y1 + heart_img.shape[0]
        self.draw_object_imgs(
            frame,
            heart_img,
            x1,
            y1,
            x2,
            y2,
            1,
        )

    def draw_bpm_text(self, frame, bpm, x, y):
        text = str(int(bpm))
        textSize = self.ft.getTextSize(text, fontHeight=17, thickness=-1)[0]
        self.ft.putText(
            img=frame,
            text=text,
            org=(x - 65 + 1, y + 10 + textSize[1] + textSize[1] // 2 + 1),
            fontHeight=17,
            color=(0, 0, 0),
            thickness=-1,
            line_type=cv2.LINE_AA,
            bottomLeftOrigin=True,
        )
        self.ft.putText(
            img=frame,
            text=text,
            org=(x - 65, y + 10 + textSize[1] + textSize[1] // 2),
            fontHeight=17,
            color=(255, 255, 255),
            thickness=-1,
            line_type=cv2.LINE_AA,
            bottomLeftOrigin=True,
        )
        self.draw_object_imgs(
            frame,
            self.bpm_img,
            x - 65 + textSize[0] - 10,
            y - 5,
            x - 65 + textSize[0] - 10 + self.text_icon_size,
            y - 5 + self.text_icon_size,
            1,
        )

    def measure_pulse(self, frame, bbox, landmarks=None):
        forehead_x = 0.5
        y_start = 0.2
        width = 0.3
        height = 0.15
        col = (100, 255, 100)

        if self.find_faces:
            # Clearing out the data buffers if not estimating an heart rate
            self.data_buffer_a = collections.deque(maxlen=self.buffer_size)
            self.data_buffer_b = collections.deque(maxlen=self.buffer_size)
            self.times_a = collections.deque(maxlen=self.buffer_size)
            self.times_b = collections.deque(maxlen=self.buffer_size)

            if bbox != []:
                self.face_rect = [
                    bbox[0],
                    bbox[1],
                    bbox[2] - bbox[0],
                    bbox[3] - bbox[1],
                ]
                if landmarks is not None:
                    left_eye_x = landmarks[2]
                    right_eye_x = landmarks[0]
                    forehead_x = right_eye_x + (left_eye_x - right_eye_x) / 2
            # BPM is set to 0 if not estimating
            return 0
        if self.face_rect is None:
            # BPM is set to 0 if not face is present
            return 0
        if bbox != []:
            self.face_rect = [
                bbox[0],
                bbox[1],
                bbox[2] - bbox[0],
                bbox[3] - bbox[1],
            ]
            if landmarks is not None:
                left_eye_x = landmarks[2]
                right_eye_x = landmarks[0]
                forehead_x = right_eye_x + (left_eye_x - right_eye_x) / 2

        # Extract the skin region around the forehead for estimation
        forehead1 = self.get_subface_coord(forehead_x, y_start, width, height)

        # Using the green channel for estimation
        _, v2, v3 = self.get_subface_means(frame, forehead1)

        # Select which data buffer to store the new measurement
        selected_buffer, selected_times = self.select_storing_buffer(
            v2, self.pixel_threshold
        )

        selected_buffer.append(v2)
        selected_times.append(time.time() - self.t0)

        # Update the data buffer with latest measurement
        self.update_storing_buffer(selected_buffer, selected_times)

        other_bpm = 0
        other_data_length = 0

        # If there is already a steady signal stored in another data buffer, obtain its bpm value for later use.
        if self.current_buffer == "a":
            other_bpm = self.calculate_bpm(self.data_buffer_b, self.times_b)
            other_data_length = len(self.data_buffer_b)
        else:
            other_bpm = self.calculate_bpm(self.data_buffer_a, self.times_a)
            other_data_length = len(self.data_buffer_a)

        total_data_length = other_data_length + len(selected_buffer)

        processed = np.array(selected_buffer)
        L = len(selected_buffer)
        # Start heartrate estimation only if there are enough measurements stored in the data buffer
        if L > self.essential_frame_accumulated:
            self.fps = float(L) / (selected_times[-1] - selected_times[0])
            even_times = np.linspace(selected_times[0], selected_times[-1], L)
            interpolated = np.interp(even_times, selected_times, processed)
            interpolated = np.hamming(L) * interpolated
            interpolated = interpolated - np.mean(interpolated)
            raw = np.fft.rfft(interpolated)
            phase = np.angle(raw)
            self.fft = np.abs(raw)
            self.freqs = float(self.fps) / L * np.arange(L / 2 + 1)
            freqs = 60.0 * self.freqs
            idx = np.where((freqs > 50) & (freqs < 220))
            pruned = self.fft[idx]
            phase = phase[idx]

            pfreq = freqs[idx]
            self.freqs = pfreq
            self.fft = pruned
            idx2 = np.argmax(pruned)

            t = (np.sin(phase[idx2]) + 1.0) / 2.0
            t = 0.9 * t + 0.1
            self.alpha = t

            gap = (self.buffer_size - L) / self.fps

            # If the selected data buffer is not completely filled, we don't consider it a steady signal.
            if gap:
                # Since the signal is not considered steady, we use a steady signal from the other data buffer to
                # perform weighting on the final result
                this_bpm = self.freqs[idx2]
                other_portion = (other_data_length / total_data_length) * other_bpm
                this_portion = (len(selected_buffer) / total_data_length) * this_bpm
                self.bpm = other_portion + this_portion
                self.bpms.append(self.bpm)
                if other_bpm == 0:
                    # If there isn't a steady signal in the other data buffer, we don't return the current value.
                    return 0
                else:
                    # If a weighted value is obtained, we return the moving average of the value.
                    return self.bpms.calc_average()
            else:
                self.bpm = self.freqs[idx2]
                self.bpms.append(self.bpm)
                return self.bpms.calc_average()
        else:
            if other_data_length >= self.buffer_size:
                return self.bpms.calc_average()
            else:
                return 0

    def draw_vitals(self, frame, bpm=0):
        # Make sure a face is detected before drawing anything
        if self.face_rect is not None:
            x, y, w, h = self.face_rect
            self.draw_disconnected_rect(frame, (x, y), (x + w, y + h), COLOR[1], 2)
            test = ""
            if not self.find_faces:
                text = "Press 'S' to stop heart rate estimation"
                if bpm == 0:
                    self.draw_fade_heart(frame, x, y, w, h)
                else:
                    self.draw_pump_heart(frame, x, y, self.alpha)
                    self.draw_bpm_text(frame, bpm, x, y)
            else:
                text = "Press 'S' to begin heart rate estimation"

            textSize = self.ft.getTextSize(text, fontHeight=17, thickness=-1)[0]
            self.ft.putText(
                img=frame,
                text=text,
                org=(10 + 1, 10 + textSize[1] + 1),
                fontHeight=17,
                color=(0, 0, 0),
                thickness=-1,
                line_type=cv2.LINE_AA,
                bottomLeftOrigin=True,
            )
            self.ft.putText(
                img=frame,
                text=text,
                org=(10, 10 + textSize[1]),
                fontHeight=17,
                color=(255, 255, 255),
                thickness=-1,
                line_type=cv2.LINE_AA,
                bottomLeftOrigin=True,
            )


# Links are represented by a 24-bit word
# when autocal reads link status from explog, we see:
# 0x<******>00<******> -- the first group is cdrStat and second is rxtStat
# 1's mean the links are up.
import re
import numpy as np

class Links:
    """ Class to deal with link errors, both regionslips and linklosses.  Be careful not to mix."""
    def __init__( self , initial=0xffffff ):
        self.bits        = 24
        self.state       = self.hex_to_array( initial )
        self.bad_links   = ~self.state
        self.fail_count  = np.zeros( self.bits , int )
        self.first_fails = 24*['']

    def add_fail( self , hex_num , flow_name ):
        # Interpret hex
        new_state  = self.hex_to_array( hex_num )
        new_bad    = ~new_state

        # update state and fail count
        self.state[new_bad]       = False
        self.bad_links            = ~self.state
        self.fail_count[new_bad] += 1

        # record new first_fail.  If this is done during active parsing, only write if it's ''.
        for idx in np.where( new_bad )[0]:
            if self.first_fails[idx] == '':
                self.first_fails[idx] = flow_name

    def show_state( self ):
        print('{}:\t{}'.format( self.array_to_hex( self.state ) , self.array_to_binary( self.state )))
        print('Bad links:\t' + ', '.join( [ str(x) for x in np.where( self.bad_links )[0] ] ) )

    def calc_metrics( self ):
        """
        Calculates metrics of interest.  They should be flattened for database migration purposes.
        autoCal plugin will leave them as nested dictionaries.  On db, we can flatten from scraperdb.models.
        """
        metrics = { 'summary_state' : self.array_to_hex( self.state ) ,
                    'summary_fails' : [ x for x in np.where( self.bad_links)[0] ] ,
                    'link'          : {} }
        for i in range(self.bits):
            metrics['link'][str(i)] = { 'first' : self.first_fails[i] , 'instances' : self.fail_count[i] }

        self.metrics = metrics

    @classmethod
    def from_misc( cls, misc_dict, regionslips=True ):
        """
        This method aims to help with parsing of misc fields in the normalduck database to
        get data on synclink fails.  Assumes misc_dict is a dict, e.g. ( json.loads( x.misc ) ).
        """
        # Filesorter blatantly stolen from Scott.
        def filesorter( filename ):
            section = '_'.join(filename.split( '_' )[:-1])
            index   = int( filename.split( '_' )[-1] )
            sections = [ 'beadfind_pre', 'prerun', 'extraG', 'acq' ]
            try:
                sid = sections.index( section )
            except:
                sid = -1
            return sid*10000+index

        rs   = re.compile( 'regionslips_([a-f0-9]*)_instances' )
        ll   = re.compile( 'linklosses_([a-f0-9]{6})00([a-f0-9]*)_instances' )
        data = []

        # Pick out fails from dictionary
        if regionslips:
            print( 'Analyzing regionslips from misc field . . .' )
            for k in misc_dict:
                m = rs.match( k )
                if m:
                    flow            = misc_dict['regionslips_{}_first'.format( m.group(1) ) ]
                    hex_count_first = ( m.group(1), int(misc_dict[k]), flow )
                    data.append( hex_count_first )
        else:
            print( 'Analyzing linklosses from misc field . . .' )
            for k in misc_dict:
                m = ll.match( k )
                if m:
                    flow            = misc_dict['linklosses_{}00{}_first'.format( m.group(1) , m.group(2) ) ]
                    hex_count_first = ( m.group(2), int(misc_dict[k]), flow )
                    data.append( hex_count_first )

        links = cls( )
        for (hex_num, count, flow_name) in data:
            print( hex_num, count, flow_name )
            # Interpret hex
            new_state  = links.hex_to_array( hex_num )
            new_bad    = ~new_state

            # update state and fail count
            links.state[new_bad]       = False
            links.bad_links            = ~links.state
            links.fail_count[new_bad] += count

            # record new first_fail.  This requires some sorting and ingenuity.
            for idx in np.where( new_bad )[0]:
                first_fail = links.first_fails[idx]
                if first_fail == '':
                    links.first_fails[idx] = flow_name
                else:
                    # We have to sort and pick earliest.
                    links.first_fails[idx] = sorted([ first_fail , flow_name ], key=filesorter )[0]

        return links

    @classmethod
    def from_nested_dict( cls, nested_dict, regionslips=True ):
        """
        This method aims to help with parsing of misc fields in the normalduck database to
        get data on synclink fails.  Assumes nested_dict is a nested dict, from a pure dictionary
        just like what would come from autoCal plugin or our explog reading code.
        """
        # Filesorter blatantly stolen from Scott.
        def filesorter( filename ):
            section = '_'.join(filename.split( '_' )[:-1])
            index   = int( filename.split( '_' )[-1] )
            sections = [ 'beadfind_pre', 'prerun', 'extraG', 'acq' ]
            try:
                sid = sections.index( section )
            except:
                sid = -1
            return sid*10000+index

        data = []

        # Pick out fails from dictionary
        if regionslips:
            print( 'Analyzing regionslips . . .' )
            rs_data = nested_dict
            for k in [key for key in rs_data if key != 'total']:
                hex_num         = str(k)
                instances       = rs_data[k]['instances']
                flow            = rs_data[k]['first']
                hex_count_first = ( hex_num, instances, flow )
                data.append( hex_count_first )
        else:
            print( 'Analyzing linklosses . . .' )
            ll_data = nested_dict
            for k in [key for key in ll_data if key != 'total']:
                hex_num         = str(k[-6:])
                instances       = ll_data[k]['instances']
                flow            = ll_data[k]['first']
                hex_count_first = ( hex_num, instances, flow )
                data.append( hex_count_first )

        links = cls( )
        for (hex_num, count, flow_name) in data:
            print( hex_num, count, flow_name )
            # Interpret hex
            new_state  = links.hex_to_array( hex_num )
            new_bad    = ~new_state

            # update state and fail count
            links.state[new_bad]       = False
            links.bad_links            = ~links.state
            links.fail_count[new_bad] += count

            # record new first_fail.  This requires some sorting and ingenuity.
            for idx in np.where( new_bad )[0]:
                first_fail = links.first_fails[idx]
                if first_fail == '':
                    links.first_fails[idx] = flow_name
                else:
                    # We have to sort and pick earliest.
                    links.first_fails[idx] = sorted([ first_fail , flow_name ], key=filesorter )[0]

        return links

    @staticmethod
    def int_to_array( num24 ):
        return np.array( [ bool(int(x)) for x in list( np.binary_repr( num24 , 24 ) ) ], dtype=bool)

    @staticmethod
    def hex_to_array( hex_num ):
        # We need the int statement to be agnostic to whether hex_num is 0xffffff or 'ffffff'
        try:
            integer = int( hex_num )
        except ValueError:
            integer = int( hex_num , 16 )
        # Note that the LSB is link 0.  This was backwards previously.
        arr = np.array( [ bool(int(x)) for x in list( np.binary_repr( integer , 24 ) ) ], dtype=bool)
        return arr[::-1]

    @staticmethod
    def array_to_binary( arr ):
        return ''.join( [ str(int(x)) for x in arr[::-1] ] )

    @staticmethod
    def array_to_int( arr ):
        return int( ''.join( [ str(int(x)) for x in arr[::-1] ] ) , 2 )

    @staticmethod
    def array_to_hex( arr ):
        return hex( int( ''.join( [ str(int(x)) for x in arr[::-1] ] ) , 2 ) )


from .base_lot import *
import numpy as np
import os
from .units import *

#TODO get rid of get_energy
class QChem(Lot):

    def run(self,geom,multiplicity):

        tempfilename = 'tempQCinp'
        tempfile = open(tempfilename,'w')
        if self.lot_inp_file == False:
            tempfile.write(' $rem\n')
            tempfile.write(' JOBTYPE FORCE\n')
            tempfile.write(' EXCHANGE {}\n'.format(self.functional))
            tempfile.write(' SCF_ALGORITHM rca_diis\n')
            tempfile.write(' SCF_MAX_CYCLES 300\n')
            tempfile.write(' BASIS {}\n'.format(self.basis))
            #tempfile.write(' ECP LANL2DZ \n')
            tempfile.write(' WAVEFUNCTION_ANALYSIS FALSE\n')
            tempfile.write(' GEOM_OPT_MAX_CYCLES 300\n')
            tempfile.write('scf_convergence 6\n')
            tempfile.write(' SYM_IGNORE TRUE\n')
            tempfile.write(' SYMMETRY   FALSE\n')
            tempfile.write('molden_format true\n')
            tempfile.write(' $end\n')
            tempfile.write('\n')
            tempfile.write('$molecule\n')
        else:
            with open(self.lot_inp_file) as lot_inp:
                lot_inp_lines = lot_inp.readlines()
            for line in lot_inp_lines:
                tempfile.write(line)

        tempfile.write('{} {}\n'.format(self.charge,multiplicity))
        if os.path.isfile("link.txt"):
            with open("link.txt") as link:
                link_lines = link.readlines()
            tmp_geom = [list(i) for i in geom]
            for i,coord in enumerate(tmp_geom):
                coord.append(link_lines[i].rstrip('\n'))
                for i in coord:
                    tempfile.write(str(i)+' ')
                tempfile.write('\n')
        else:
            for coord in geom:
                for i in coord:
                    tempfile.write(str(i)+' ')
                tempfile.write('\n')
        tempfile.write('$end')
        tempfile.close()

        cmd = "qchem -nt {} -save {} {}.qchem.out {}.{}".format(self.nproc,tempfilename,tempfilename,self.node_id,multiplicity)

        os.system(cmd)

        efilepath = os.environ['QCSCRATCH']
        efilepath += '/{}.{}/GRAD'.format(self.node_id,multiplicity)
        with open(efilepath) as efile:
            elines = efile.readlines()

        temp = 0
        for lines in elines:
            if temp == 1:
                self.E.append((multiplicity,float(lines.split()[0])))
                break
            if "$" in lines:
                temp += 1

        gradfilepath = os.environ['QCSCRATCH']
        gradfilepath += '/{}.{}/GRAD'.format(self.node_id,multiplicity)

        with open(gradfilepath) as gradfile:
            gradlines = gradfile.readlines()
        temp = 0
        tmp=[]
        for lines in gradlines:
            if '$' in lines:
                temp+=1
            elif temp == 2:
                tmpline = lines.split()
                tmp.append([float(i) for i in tmpline])
            elif temp == 3:
                break
        self.grada.append((multiplicity,tmp))

        return

    def get_energy(self,coords,multiplicity,state):
        #if self.has_nelectrons==False:
        #    for i in self.states:
        #        self.get_nelec(geom,i[0])
        #    self.has_nelectrons==True
        if self.hasRanForCurrentCoords==False or (coords != self.currentCoords).any():
            self.currentCoords = coords.copy()
            geom = manage_xyz.np_to_xyz(self.geom,self.currentCoords)
            self.runall(geom)
            self.hasRanForCurrentCoords=True
        tmp = self.search_tuple(self.E,multiplicity)
        return np.asarray(tmp[state][1])*KCAL_MOL_PER_AU

    def get_gradient(self,coords,multiplicity,state):
        #if self.has_nelectrons==False:
        #    for i in self.states:
        #        self.get_nelec(geom,i[0])
        #    self.has_nelectrons==True
        if self.hasRanForCurrentCoords==False or (coords != self.currentCoords).any():
            self.currentCoords = coords.copy()
            geom = manage_xyz.np_to_xyz(self.geom,self.currentCoords)
            self.runall(geom)
        tmp = self.search_tuple(self.grada,multiplicity)
        return np.asarray(tmp[state][1])*ANGSTROM_TO_AU

    @classmethod
    def copy(cls,lot,**kwargs):
        base = os.environ['QCSCRATCH']
        for state in self.states:
            multiplicity = state[0]
            efilepath_old=base+ '/{}.{}'.format(self.node_id,multiplicity)
            efilepath_new =base+ '/{}.{}'.format(node_id,multiplicity)
            os.system('cp -r ' + efilepath_old +' ' + efilepath_new)
        return cls(lot.options.copy().set_values(options))


# plot the old and new curves for the same sightlines
import argparse
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import copy

from measure_extinction.extdata import ExtData

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--png", help="save figure as a png file", action="store_true")
    parser.add_argument("--pdf", help="save figure as a pdf file", action="store_true")
    args = parser.parse_args()

    # plotting setup for easier to read plots
    fontsize = 14
    font = {"size": fontsize}
    mpl.rc("font", **font)
    mpl.rc("lines", linewidth=1)
    mpl.rc("axes", linewidth=2)
    mpl.rc("xtick.major", width=2)
    mpl.rc("xtick.minor", width=2)
    mpl.rc("ytick.major", width=2)
    mpl.rc("ytick.minor", width=2)

    fig, fax = plt.subplots(nrows=2, ncols=3, figsize=(16, 8), sharey=True, sharex=True)

    ax = [fax[0, 0], fax[0, 1], fax[1, 0], fax[1, 1], fax[0, 2]]

    ptypes = ["IUE", "BAND"]
    psym = ["-", "o"]

    snames = ["AzV18", "AzV23", "AzV214", "AzV398", "AzV456"]
    ofiles = [
        "azv18_azv462_ext.fits",
        "azv23_azv404_ext.fits",
        "azv214_azv380_ext.fits",
        "azv398_azv289_ext.fits",
        "azv456_azv70_ext.fits",
    ]
    nfiles = [
        "azv_018_ext.fits",
        "azv_023_ext.fits",
        "azv_214_ext.fits",
        "azv_398_ext.fits",
        "azv_456_ext.fits",
    ]
    pcol = ["r", "b", "g", "c", "m"]
    for k, ofile in enumerate(ofiles):

        oext = ExtData(filename=f"prev/{ofile}")
        next = ExtData(filename=f"fits/{nfiles[k]}")

        dext = copy.deepcopy(oext)

        for curtype, csym in zip(ptypes, psym):
            oext.exts[curtype][oext.npts[curtype] == 0] = np.nan
            if curtype == "IUE":
                olabel = "G03"
                nlabel = "this work"
                dlabel = "this work - G03"
                dlabel2 = f"{snames[k]} diff"
            else:
                olabel = None
                nlabel = None
                dlabel = None
                dlabel2 = None
            ax[k].plot(
                1.0 / oext.waves[curtype], oext.exts[curtype], f"g{csym}", label=olabel
            )
            next.exts[curtype][next.npts[curtype] == 0] = np.nan
            ax[k].plot(
                1.0 / next.waves[curtype],
                next.exts[curtype],
                f"b{csym}",
                label=nlabel,
            )

            print(len(oext.waves[curtype]), len(next.waves[curtype]))
            if curtype == "BAND":  # need to remove RI measurements as not in new ext
                gphot = (oext.waves[curtype].value <= 0.6) | (oext.waves[curtype].value >= 1.0)

                dext.waves[curtype] = next.waves[curtype]
                dext.exts[curtype] = next.exts[curtype] - oext.exts[curtype][gphot]
                # dext.name[curtype] = next.name[curtype]
                dext.npts[curtype] = next.npts[curtype]
                dext.uncs[curtype] = next.uncs[curtype]
            else:
                inew = np.interp(
                    oext.waves[curtype], next.waves[curtype], next.exts[curtype]
                )
                dext.exts[curtype] = inew - oext.exts[curtype]

            ax[k].plot(
                1.0 / dext.waves[curtype],
                dext.exts[curtype],
                f"r{csym}",
                label=dlabel,
            )
            fax[1, 2].plot(
                1.0 / dext.waves[curtype],
                dext.exts[curtype],
                f"{pcol[k]}-",
                label=dlabel2,
            )

        ax[k].set_title(snames[k])

        # save differential extinction
        tfile = nfiles[k].replace(".fits", "_diff.fits")
        dext.save(f"fits/{tfile}")

    ax[0].legend()
    fax[1, 2].legend()

    ax[0].set_ylim(-2.0, 5.0)

    fax[1, 0].set_xlabel(r"$\lambda$ [$\mu$m$^{-1}$]")
    fax[1, 1].set_xlabel(r"$\lambda$ [$\mu$m$^{-1}$]")
    fax[1, 2].set_xlabel(r"$\lambda$ [$\mu$m$^{-1}$]")

    fax[0, 0].set_ylabel(r"$E(\lambda - V)$")
    fax[1, 0].set_ylabel(r"$E(\lambda - V)$")

    fig.tight_layout()

    # plot or save to a file
    outname = "old_vs_new"
    if args.png:
        fig.savefig(outname + ".png")
    elif args.pdf:
        fig.savefig(outname + ".pdf")
    else:
        plt.show()


import pygame
from pygame.locals import *
from OpenGL.GL import *
from OpenGL.GLU import *
from sensor import Sensor
from madgwick import Madgwick
import numpy as np
import math

vertices = (
    (1, -1, -1),
    (1, 1, -1),
    (-1, 1, -1),
    (-1, -1, -1),
    (1, -1, 1),
    (1, 1, 1),
    (-1, -1, 1),
    (-1, 1, 1)
    )

edges = (
    (0,1),
    (0,3),
    (0,4),
    (2,1),
    (2,3),
    (2,7),
    (6,3),
    (6,4),
    (6,7),
    (5,1),
    (5,4),
    (5,7)
    )

def Cube():
    glBegin(GL_LINES)
    for edge in edges:
        for vertex in edge:
            glVertex3fv(vertices[vertex])
    glEnd()

def main():
    pygame.init()
    display = (800,600)
    pygame.display.set_mode(display, DOUBLEBUF|OPENGL)

    gluPerspective(45, (display[0]/display[1]), 0.1, 50.0)
    glTranslatef(0.0,0.0, -5)

    mpu9250 = Sensor('/dev/ttyACM0', 115200)
    m = Madgwick()

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                quit()
        # sensor data
        # mpu9250.update_values()

        # using average values
        mpu9250.update_values_with_moving_average(15)
        accel = mpu9250._accel_avg
        gyro = mpu9250._gyro_avg
        mag = mpu9250._mag_avg
        # madgwick algorithm
        m.update(accel, gyro * np.pi / 180, mag)
        q = m._quaternion_avg
        q = np.squeeze(np.asarray(q))

        glMatrixMode(GL_MODELVIEW)
        glLoadMatrixf(Madgwick.quat_to_rotmatrix(q))
        glScalef(0.2, 0.2, 0.2)
        #glRotatef(q[0], q[1], q[2], q[3])
        #glRotatef(1, 3, 1, 1)
        glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT)
        Cube()
        pygame.display.flip()
        pygame.time.wait(10)

if __name__ == "__main__":
    main()


import os
import os.path as osp

import deepdish as dd
import numpy as np
import tensorflow as tf

from src.models import (
    batch_pred_omega,
    get_image_encoder,
    get_prediction_model,
    get_temporal_encoder,
)
from src.omega import (
    OmegasPred,
)
from src.tf_smpl.batch_smpl import SMPL

class TesterPred(object):

    def __init__(self, config, sequence_length, resnet_path='', sess=None,
                 precomputed_phi=False):
        self.config = config
        self.load_path = config.load_path
        tf.set_random_seed(config.seed)

        self.num_conv_layers = 3
        self.fov = self.num_conv_layers * 4 + 1
        self.sequence_length = sequence_length
        self.use_delta_from_pred = config.use_delta_from_pred
        self.use_optcam = config.use_optcam
        self.precomputed_phi = precomputed_phi

        # Config + path
        if not config.load_path:
            raise Exception(
                'You need to specify `load_path` to load a pretrained model'
            )
        if not osp.exists(config.load_path + '.index'):
            print('{} doesnt exist'.format(config.load_path))
            import ipdb
            ipdb.set_trace()

        # Data
        self.batch_size = config.batch_size
        self.img_size = config.img_size
        self.E_var = []
        self.pad_edges = config.pad_edges

        self.smpl_model_path = config.smpl_model_path
        self.use_hmr_ief = False

        self.num_output = 85

        if precomputed_phi:
            input_size = (self.batch_size, self.sequence_length, 2048)
        else:
            input_size = (self.batch_size, self.sequence_length,
                          self.img_size, self.img_size, 3)
        self.images_pl = tf.placeholder(tf.float32, shape=input_size)

        strip_size = (self.batch_size, self.fov, 2048)
        self.movie_strips_pl = tf.placeholder(tf.float32, shape=strip_size)

        # Model Spec
        self.f_image_enc = get_image_encoder()
        self.f_temporal_enc = get_temporal_encoder()
        self.f_prediction_ar = get_prediction_model()

        self.smpl = SMPL(self.smpl_model_path)
        self.omegas_movie_strip = self.make_omega_pred()
        self.omegas_pred = self.make_omega_pred(use_optcam=True)

        # HMR Model Params
        self.num_stage = 3
        self.total_params = 85

        self.load_mean_omega()
        self.build_temporal_encoder_model()
        self.build_auto_regressive_model()
        self.update_E_vars()

        if sess is None:
            options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)
            self.sess = tf.Session(config=tf.ConfigProto(gpu_options=options))
        else:
            self.sess = sess

        # Load data.
        self.prepare(resnet_path)

    def make_omega_pred(self, use_optcam=False):
        return OmegasPred(
            config=self.config,
            smpl=self.smpl,
            use_optcam=use_optcam,
            vis_max_batch=self.batch_size,
            is_training=False,
        )

    def update_E_vars(self):
        trainable_vars = tf.contrib.framework.get_variables()
        trainable_vars_e = [var for var in trainable_vars
                            if var.name[:2] != 'D_']
        self.E_var.extend(trainable_vars_e)

    def load_mean_omega(self):
        # Initialize scale at 0.9
        mean_path = os.path.join(os.path.dirname(self.smpl_model_path),
                                 'neutral_smpl_meanwjoints.h5')
        mean_vals = dd.io.load(mean_path)

        mean_cams = [0.9, 0, 0]
        # 72D
        mean_pose = mean_vals['pose']
        mean_pose[:3] = 0.
        mean_pose[0] = np.pi
        # 10D
        mean_shape = mean_vals['shape']

        mean_vals = np.hstack((mean_cams, mean_pose, mean_shape))
        # Needs to be 1 x 85
        mean_vals = np.expand_dims(mean_vals, 0)
        self.mean_var = tf.Variable(
            mean_vals,
            name='mean_param',
            dtype=tf.float32,
            trainable=True
        )
        mean_cams = self.mean_var[0, :3]
        mean_pose = self.mean_var[0, 3:3+72]
        mean_shape = self.mean_var[0, 3+72:]

        self.mean_vars = [mean_cams, mean_pose, mean_shape]
        mean_cams = tf.tile(tf.reshape(mean_cams, (1,-1)),
                            (self.batch_size, 1))
        mean_shape = tf.tile(tf.reshape(mean_shape, (1,-1)),
                             (self.batch_size, 1))
        mean_pose = tf.tile(tf.reshape(mean_pose, (1, 24, 3)),
                            (self.batch_size, 1, 1))
        _, mean_joints3d, mean_poses_rot = self.smpl(
            mean_shape, mean_pose, get_skin=True)

        # Starting point for IEF.
        self.theta_mean = tf.concat((
                mean_cams,
                tf.reshape(mean_pose, (-1, 72)),
                mean_shape
        ), axis=1)

    def prepare(self, resnet_path=''):
        """
        Restores variables from checkpoint.

        Args:
            resnet_path (str): Optional path to load resnet weights.
        """
        if resnet_path and not self.precomputed_phi:
            print('Restoring resnet vars from', resnet_path)
            resnet_vars = []
            e_vars = []
            for var in self.E_var:
                if 'resnet' in var.name:
                    resnet_vars.append(var)
                else:
                    e_vars.append(var)
            resnet_saver = tf.train.Saver(resnet_vars)
            resnet_saver.restore(self.sess, resnet_path)
        else:
            e_vars = self.E_var
        print('Restoring checkpoint ', self.load_path)

        saver = tf.train.Saver(e_vars)
        saver.restore(self.sess, self.load_path)
        self.sess.run(self.mean_vars)

    def build_temporal_encoder_model(self):
        B, T = self.batch_size, self.sequence_length
        if self.precomputed_phi:
            print('loading pre-computed phi!')
            self.img_feat_full = self.images_pl
        else:
            print('Getting all image features...')
            I_t = tf.reshape(
                self.images_pl,
                (B * T, self.img_size, self.img_size, 3)
            )
            img_feat, phi_var_scope = self.f_image_enc(
                I_t,
                is_training=False,
                reuse=False,
            )
            self.img_feat_full = tf.reshape(img_feat, (B, T, -1))

        omega_mean = tf.tile(self.theta_mean, (self.sequence_length, 1))

        # At training time, we only use first 40. Want to make sure GN
        # statistics are right.
        self.movie_strips_cond = self.f_temporal_enc(
            net=self.img_feat_full[:, :40],
            num_conv_layers=self.num_conv_layers,
            prefix='',
            reuse=None,
        )

        self.movie_strips = self.f_temporal_enc(
            net=self.img_feat_full,
            num_conv_layers=self.num_conv_layers,
            prefix='',
            reuse=True,
        )

        omega_movie_strip, _ = batch_pred_omega(
            input_features=self.movie_strips,
            batch_size=B,
            sequence_length=T,
            num_output=self.num_output,
            is_training=False,
            omega_mean=omega_mean,
            scope='single_view_ief',
            use_delta_from_pred=self.use_delta_from_pred,
            use_optcam=self.use_optcam,
        )

        self.omegas_movie_strip.append_batched(omega_movie_strip)
        self.omegas_movie_strip.compute_smpl()

    def build_auto_regressive_model(self):
        omega_mean = tf.tile(self.theta_mean, (self.fov, 1))
        input_movie_strips = self.movie_strips_pl  # B x 13 x 2048
        movie_strip_pred = self.f_prediction_ar(
            net=input_movie_strips,
            num_conv_layers=self.num_conv_layers,
            prefix='pred_',
            reuse=None,
        )
        omega_pred, _ = batch_pred_omega(
            input_features=movie_strip_pred,
            batch_size=self.batch_size,
            is_training=False,
            num_output=self.num_output,
            omega_mean=omega_mean,
            sequence_length=self.fov,
            scope='single_view_ief',
            predict_delta_keys=(),
            use_delta_from_pred=self.use_delta_from_pred,
            use_optcam=self.use_optcam,
        )
        # Only want the last entry.
        self.movie_strip_pred = movie_strip_pred[:, -1:]
        self.omegas_pred.append_batched(omega_pred[:, -1:])
        self.omegas_pred.compute_smpl()

    def make_fetch_dict(self, omegas, suffix=''):
        return {
            # Predictions.
            'cams' + suffix: omegas.get_cams(),
            'joints' + suffix: omegas.get_joints(),
            'kps' + suffix: omegas.get_kps(),
            'poses' + suffix: omegas.get_poses_rot(),
            'shapes' + suffix: omegas.get_shapes(),
            'verts' + suffix: omegas.get_verts(),
            'omegas' + suffix: omegas.get_raw(),
        }

    def predict_movie_strips(self, images, get_smpl=False):
        """
        Converts images to movie strip representation. Number of images should
        be equal to sequence_length. If precomputed_phi, images should be phis.

        Args:
            images (B x (2*fov-1) x H x W x 3) or (B x (2*fov-1) x 2048).
            get_smpl (bool): If True, returns all the smpl stuff.

        Returns:
            Movie strips (B x (2*fov-1) x 2048).
        """
        feed_dict = {
            self.images_pl: images,
        }
        fetch_dict = {
            'movie_strips': self.movie_strips,
            'movie_strips_cond': self.movie_strips_cond,
        }
        if get_smpl:
            fetch_dict.update(self.make_fetch_dict(self.omegas_movie_strip))
        return self.sess.run(fetch_dict, feed_dict)

    def predict_auto_regressive(self, movie_strips):
        """
        Predicts the next time step in an auto-regressive manner.

        Args:
            movie_strips (B x fov x 2048).

        Returns:

        """
        feed_dict = {
            self.movie_strips_pl: movie_strips,
        }

        fetch_dict = {
            'movie_strip': self.movie_strip_pred,
        }
        fetch_dict.update(self.make_fetch_dict(self.omegas_pred))
        result = self.sess.run(fetch_dict, feed_dict)
        return result


# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""
import cv2
import copy
import logging
from collections import OrderedDict
from sklearn import metrics
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from projects.InterpretationReID.interpretationreid.utils.logger import setup_logger
from fastreid.utils.file_io import PathManager
from .evaluator import DatasetEvaluator
from .query_expansion import aqe
from .rank import evaluate_rank
from .rerank import re_ranking
from .roc import evaluate_roc
from fastreid.utils import comm
import os
import random
from multiprocessing import Pool,Process
import time
logger = logging.getLogger(__name__)

class ReidEvaluator(DatasetEvaluator):
    def __init__(self, cfg, num_query, output_dir=None):
        self.cfg = cfg
        self._num_query = num_query
        self._output_dir = output_dir

        self.features = []
        self.pids = []
        self.camids = []

        # for visualize
        self.imgs = []
        self.imgs_path = []
        self.real_attributes = []
        self.fake_attributes = []
        self.key_attribute = None
        self.att_list = None
        self.feature_mask = []
        self.att_prec = []

        output_dir = self.cfg.OUTPUT_DIR
        if comm.is_main_process() and output_dir:
            PathManager.mkdirs(output_dir)

        rank = comm.get_rank()
        self.logger = setup_logger(output_dir, distributed_rank=rank, name="projects")

    def reset(self):
        self.features = []
        self.pids = []
        self.camids = []

        # for visualize
        self.imgs = []
        self.imgs_path = []
        self.real_attributes = []
        self.fake_attributes = []
        self.key_attribute = None
        self.att_list = None
        self.feature_mask = []
        self.att_prec = []

    def process(self, inputs, outputs):
        self.pids.extend(inputs["targets"])
        self.camids.extend(inputs["camids"])
        self.features.append(outputs.cpu())

    def name_of_attribute(self, key_attribute):
        self.key_attribute =  key_attribute

    def process_for_visualize(self, inputs, outs):
        """
        inputs:
        {
            "images": img,
            "targets": pid,
            "camids": camid,
            "img_paths": img_path,
            "img_attributes":p_attribute
        }

        outs:
        {
        "outputs": outputs,  -> tensor : n x 2048
        "att_list": att_list, -> [tensor, tensor, ... tensor ] : n x 2048
        "feature_mask":feature_mask, -> tensor : n x 23 x 7 x 7
        "fake_attributes": fake_attributes,
        "real_attributes": real_attributes

        }

        """

        if self.cfg.VISUAL.OPEN:
            self.imgs.append(inputs["images"].cpu())
        self.imgs_path.extend(inputs["img_paths"])
        self.real_attributes.append(outs["real_attributes"].cpu())
        #self.fake_attributes.append(outs["fake_attributes"].cpu())
        self.feature_mask.append(outs["feature_mask"].cpu())
        if True:
            cls_outputs = outs['att_heads']['cls_outputs'].clone().detach()
            self.att_prec.append(torch.where(cls_outputs>0.5,torch.ones_like(cls_outputs),torch.zeros_like(cls_outputs)).cpu())

        if self.att_list == None:
            self.att_list = outs["att_list"]
            self.len_att_list = len(self.att_list)
        else:
            for i in range(self.len_att_list ):
                self.att_list[i] = torch.cat([self.att_list[i].cpu(),outs["att_list"][i].cpu()],dim=0)

    @staticmethod
    def cal_dist(metric: str, query_feat: torch.tensor, gallery_feat: torch.tensor , out_is_torch=False):
        assert metric in ["cosine", "euclidean"], "must choose from [cosine, euclidean], but got {}".format(metric)
        if metric == "cosine":
            dist = 1 - torch.mm(query_feat, gallery_feat.t())
        else:
            m, n = query_feat.size(0), gallery_feat.size(0)
            xx = torch.pow(query_feat, 2).sum(1, keepdim=True).expand(m, n)
            yy = torch.pow(gallery_feat, 2).sum(1, keepdim=True).expand(n, m).t()
            dist = xx + yy
            dist.addmm_(query_feat, gallery_feat.t(), beta=1, alpha=-2)
            dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability
        if out_is_torch:
            return dist.cpu() #torch.tensor
        else:
            return dist.cpu().numpy()
#TODO CXD Visualizer

    def explain_eval(self,dist_fake_stack, query_attribute, gallery_attribute, lamda=2.0):

        _ , num_att = query_attribute.size()
        num_att = float(num_att)
        n, m, _ = dist_fake_stack.size()

        dist_real_stack = query_attribute.unsqueeze(1) * gallery_attribute.unsqueeze(
            0)  # -1 means different, 1 means same

        assert dist_real_stack.size() == dist_fake_stack.size()
        dist_real_stack_01 = (1.0 - dist_real_stack) / 2.0  # 1 means different, 0 means same
        dist_real_stack_01_reverse = 1.0 - dist_real_stack_01  # 0 means different, 1 means same

        dist_num_different = dist_real_stack_01.sum(-1)  # n x m
        dist_num_same = dist_real_stack_01_reverse.sum(-1)  # n x m

        # dist_fake_stack belong to 0.0~1.0
        dist_fake_stack = dist_fake_stack / dist_fake_stack.sum(-1).unsqueeze(-1)  # n x m x num_att

        # calcuate Precsion of different attribute
        dist_fake_att_different = dist_fake_stack.argsort(dim=-1, descending=True)
        dist_fake_att = (dist_fake_att_different.argsort(dim=-1, descending=False) < dist_num_different.unsqueeze(
            -1)).float()  # n x m x num_att

        Precsion_different_allatt = (dist_fake_att * dist_real_stack_01).sum(0).sum(0)/dist_real_stack_01.sum(0).sum(0)
        self._results["diff_0"] = float(0.0)
        for i in range(Precsion_different_allatt.size(0)):
            self._results["diff_"+str(i+1)+"_"+self.key_attribute[i]] = float(Precsion_different_allatt[i])

        Precsion_different = (dist_fake_att * dist_real_stack_01).sum(-1) / dist_num_different.clamp(min=1.0)
        Precsion_different = Precsion_different.sum() / (n * m - (dist_num_different == 0).float().sum()).clamp(min=1.0)
        self._results["Ex_Precsion_d"] = float(Precsion_different)
        # print(dist_fake_att_different)

        # calcuate Precsion of same attribute
        dist_same_equal_numatt = (dist_num_same.unsqueeze(-1) == num_att).float()  # n x m x 1
        dist_overflow = (dist_fake_stack > lamda / num_att).float() * dist_same_equal_numatt
        dist_fake_att_same = dist_fake_stack.argsort(dim=-1, descending=False)
        dist_fake_att = (dist_fake_att_same.argsort(dim=-1, descending=False) < dist_num_same.unsqueeze(
            -1)).float() - dist_overflow  # n x m x num_att

        Precsion_same_allatt = (dist_fake_att * dist_real_stack_01_reverse).sum(0).sum(0)/dist_real_stack_01_reverse.sum(0).sum(0)
        self._results["same_0"] = float(0.0)
        for i in range(Precsion_different_allatt.size(0)):
            self._results["same_"+str(i+1)+"_"+self.key_attribute[i]] = float(Precsion_same_allatt[i])

        Precsion_same = (dist_fake_att * dist_real_stack_01_reverse).sum(-1) / dist_num_same.clamp(min=1.0)
        Precsion_same = Precsion_same.sum() / (n * m - (dist_num_same == 0).float().sum()).clamp(min=1.0)
        self._results["Ex_Precsion_s"] = float(Precsion_same)

        # print(dist_fake_att_different[0,0])
        raw_cmc = dist_real_stack_01[torch.tensor(range(n)).reshape(-1, 1, 1), torch.tensor(range(m)).reshape(1, -1,
                                                                                1), dist_fake_att_different.reshape(
            n, m, int(num_att))]
        tmp_cmc = raw_cmc.cumsum(dim=-1)
        div_cmc = torch.tensor(range(1, int(num_att) + 1)).repeat(n, m, 1)
        tmp_cmc = tmp_cmc / div_cmc * raw_cmc
        AP_d = tmp_cmc.sum(dim=-1) / dist_num_different.clamp(min=1)
        # all_AP_d.append(AP.sum())

        dist_same_equal_numatt = -2 * (dist_same_equal_numatt - 0.5)  # -1 means same attribute == num_att
        dist_fake_att_same = (dist_fake_stack * dist_same_equal_numatt).argsort(dim=-1, descending=False)

        raw_cmc = (dist_real_stack_01_reverse - dist_overflow)[
            torch.tensor(range(n)).reshape(-1, 1, 1), torch.tensor(range(m)).reshape(1, -1,
                                                                                1), dist_fake_att_same.reshape(n,
                                                                                m,
                                                                                int(
                                                                                num_att))]
        tmp_cmc = raw_cmc.cumsum(dim=-1)
        div_cmc = torch.tensor(range(1, int(num_att) + 1)).repeat(n, m, 1)
        tmp_cmc = tmp_cmc / div_cmc * raw_cmc
        AP_s = tmp_cmc.sum(dim=-1) / dist_num_same.clamp(min=1)

        self._results["Ex_mAP_d"] = AP_d.sum() / (n * m - (dist_num_different == 0).float().sum()).clamp(min=1.0)
        self._results["Ex_mAP_s"] = AP_s.sum() / (n * m - (dist_num_same == 0).float().sum()).clamp(min=1.0)

        return copy.deepcopy(self._results)

    def visualize(self):
        if comm.get_world_size() > 1:
            assert False
            #TODO CXD
        else:
            features = self.features

            # pids = self.pids
            # camids = self.camids

            # for visualize
            if self.cfg.VISUAL.OPEN:
                imgs_path = self.imgs_path
                imgs = self.imgs
            real_attributes = self.real_attributes
            fake_attributes = self.fake_attributes
            key_attribute = self.key_attribute
            att_list = self.att_list
            feature_mask = self.feature_mask

        #att_names: [str,...,str]
        choose_att_names = key_attribute
        if self.cfg.VISUAL.OPEN:
            imgs = torch.cat(imgs, dim=0)
        features = torch.cat(features, dim=0)
        feature_mask = torch.cat(feature_mask,dim=0)
        real_attributes = torch.cat(real_attributes, dim=0)
        #fake_attributes = torch.cat(fake_attributes, dim=0)
        fake_attributes = torch.zeros_like(real_attributes) # adapt
        if self.cfg.VISUAL.OPEN:
            _, _, h_imgs, w_imgs = imgs.size()
        else:
            h_imgs, w_imgs = 384,192

        # query feature, person ids and camera ids
        query_features = features[:self._num_query]
        # query_pids = np.asarray(pids[:self._num_query])
        # query_camids = np.asarray(camids[:self._num_query])
        if self.cfg.VISUAL.OPEN:
            query_imgs_path  = imgs_path[:self._num_query]
            query_imgs = imgs[:self._num_query]
        query_real_attributes = real_attributes[:self._num_query]
        query_fake_attributes = fake_attributes[:self._num_query]
        query_feature_mask = feature_mask[:self._num_query]
        query_att_list = []
        for i in range(self.len_att_list):
            query_att_list.append(att_list[i][:self._num_query])

        # gallery features, person ids and camera ids
        gallery_features = features[self._num_query:]
        # gallery_pids = np.asarray(pids[self._num_query:])
        # gallery_camids = np.asarray(camids[self._num_query:])
        if self.cfg.VISUAL.OPEN:
            gallery_imgs_path = imgs_path[self._num_query:]
            gallery_imgs = imgs[self._num_query:]
        gallery_real_attributes = real_attributes[self._num_query:]
        gallery_fake_attributes = fake_attributes[self._num_query:]
        gallery_feature_mask = feature_mask[self._num_query:]
        gallery_att_list = []
        for i in range(self.len_att_list):
            gallery_att_list.append(att_list[i][self._num_query:])

        # Initialize
        if self.cfg.TEST.AQE.ENABLED:
            self.logger.info("Test with AQE setting")
            qe_time = self.cfg.TEST.AQE.QE_TIME
            qe_k = self.cfg.TEST.AQE.QE_K
            alpha = self.cfg.TEST.AQE.ALPHA
            query_features, gallery_features = aqe(query_features, gallery_features, qe_time, qe_k, alpha)
            for i in range(self.len_att_list):
                query_att_list[i],gallery_att_list[i] = aqe(query_att_list[i] , gallery_att_list[i], qe_time, qe_k, alpha)

        if self.cfg.TEST.METRIC == "cosine":
            query_features = F.normalize(query_features, dim=1)
            gallery_features = F.normalize(gallery_features, dim=1)
            for i in range(self.len_att_list):
                query_att_list[i] = F.normalize(query_att_list[i], dim=1)
                gallery_att_list[i] = F.normalize(gallery_att_list[i], dim=1)

        dist = self.cal_dist(self.cfg.TEST.METRIC, query_features, gallery_features,out_is_torch=True) # n (query) x m
        n,m = dist.size()
        dist_list = [] # [tensor ,..., tensor] -> tensor -> n x m
        for i in range(self.len_att_list):
            dist_list.append(self.cal_dist(self.cfg.TEST.METRIC, query_att_list[i] , gallery_att_list[i],out_is_torch=True))

        #calculate
        dist_list_stack = torch.stack(dist_list,dim=-1) # n x m x NUM_ATT

        #TODO CXD EXP
        fake_gap = (query_fake_attributes.unsqueeze(1)-gallery_fake_attributes.unsqueeze(0)).abs().exp()  # n x m x NUM_ATT
        loss_interpretation_withoutmean = (
                    dist - (dist_list_stack * fake_gap ).mean(dim=-1)).abs()  # n x m
        gap_real_fake = (loss_interpretation_withoutmean / dist).mean() * 100

        self.logger.info("dist:\n {}".format(dist))
        self.logger.info("fake_gap \n {}".format(fake_gap))
        self.logger.info("dist_list_stack \n {}".format(dist_list_stack))

        self.logger.info("gap between real and fake y is {} %".format(gap_real_fake))
        #print("gap between real and fake y is {} %".format(gap_real_fake))

        visiual_average_att(choose_att_names, feature_mask,real_attributes, img_size=(384, 192), output_dir=self.cfg.OUTPUT_DIR,positive=True)
        visiual_average_att(choose_att_names, feature_mask, real_attributes, img_size=(384, 192), output_dir=self.cfg.OUTPUT_DIR, positive=False)
        #print(choose_att_names)
        if self.cfg.VISUAL.OPEN:
            for i_4 in list([2,4,6,8,10,20]): # 2,4,6,8,10,20
            #for i_4 in list([0]):

                self.mkdir_id = str(time.time()).replace(".","_")
                os.mkdir(os.path.join(self.cfg.OUTPUT_DIR, "output_imgs_"+self.mkdir_id))

                for id_query in range(n):
                    if id_query % self.cfg.VISUAL.GAP_QUERY != 0:
                        continue
                # print(query_imgs_path)
                # for id_query in range(2):
                #     id_gallery_choose = gallery_imgs_path.index(list_gallery[id_query])
                #
                #     id_query = query_imgs_path.index(list_query[id_query])

                    _,indices_m = dist[id_query].sort()

                    indices_m = indices_m.numpy().tolist()

                    # type 3
                    list_choosen = [id_gallery_choose, indices_m[0], indices_m[1], indices_m[2]]

                    #path_names:[str,...,str]
                    choose_path_names = [query_imgs_path[id_query].split("/")[-1]]
                    for id_gallery in list_choosen:
                        choose_path_names.append(gallery_imgs_path[id_gallery].split("/")[-1])
                    #imgs: tensor -> n x 3 x 224 x 224
                    choose_imgs = query_imgs[id_query:id_query+1]
                    choose_imgs = torch.cat([choose_imgs,gallery_imgs[list_choosen]],dim=0)

                    #att_mask: tensor -> n x 23 x 224 x 224
                    choose_att_mask = query_feature_mask[id_query:id_query+1]
                    choose_att_mask = torch.cat([choose_att_mask,gallery_feature_mask[list_choosen]],dim=0)

                    #real_attribute: tensor -> 23 x n
                    choose_real_attributes = query_real_attributes[id_query:id_query+1]
                    choose_real_attributes = torch.cat([choose_real_attributes,gallery_real_attributes[list_choosen]],dim=0)
                    choose_real_attributes = choose_real_attributes.t()

                    #fake_attribute: tensor -> 23 x n
                    choose_fake_attributes = query_fake_attributes[id_query:id_query+1]
                    choose_fake_attributes = torch.cat([choose_fake_attributes,gallery_fake_attributes[list_choosen]],dim=0)
                    choose_fake_attributes = choose_fake_attributes.t()

                    #feature_distance: tensor  -> (1 + 23 + 1) x 4
                    choose_feature_distance = dist[id_query:id_query+1,list_choosen] #  1x4

                    for i in range(self.len_att_list):
                        #TODO CXD EXP
                        fake_feature_gap = (choose_fake_attributes[i:i+1,0:1]-choose_fake_attributes[i:i+1,1:]).abs().exp() # 1 x 4
                        fake_feature_distance = dist_list[i][id_query:id_query + 1, list_choosen] * fake_feature_gap  #1 x 4
                        choose_feature_distance = torch.cat([choose_feature_distance,fake_feature_distance],dim=0)

                    fake_sum_feature_distance = choose_feature_distance[1:].mean(dim=0).reshape(1,4)
                    choose_feature_distance = torch.cat( [choose_feature_distance,fake_sum_feature_distance ], dim=0) # 25 x 4

                    visualization_savefig(choose_path_names, choose_att_names, choose_imgs / 255.0, choose_att_mask,
                                          choose_feature_distance, choose_real_attributes, choose_fake_attributes,
                                          output_dir=os.path.join(self.cfg.OUTPUT_DIR, "output_imgs_" + self.mkdir_id),only_hot_map=False)
        #To save memory

        del self.imgs,self.fake_attributes,self.feature_mask,self.att_list
        self.imgs = []
        self.fake_attributes = []
        self.att_list = None
        self.feature_mask = []

        return dist_list_stack,query_real_attributes, gallery_real_attributes # n x m x NUM_ATT

    def evaluate_dist(self,bias_dist,q_att=None,p_att=None,query_real_attributes=None):
        if comm.get_world_size() > 1:
            comm.synchronize()
            features = comm.gather(self.features)
            features = sum(features, [])

            pids = comm.gather(self.pids)
            pids = sum(pids, [])

            camids = comm.gather(self.camids)
            camids = sum(camids, [])

            if not comm.is_main_process():
                return {}

        else:
            features = self.features
            pids = self.pids
            camids = self.camids

        features = torch.cat(features, dim=0)
        # query feature, person ids and camera ids
        query_features = features[:self._num_query]
        query_pids = np.asarray(pids[:self._num_query])
        query_camids = np.asarray(camids[:self._num_query])

        # gallery features, person ids and camera ids
        gallery_features = features[self._num_query:]
        gallery_pids = np.asarray(pids[self._num_query:])
        gallery_camids = np.asarray(camids[self._num_query:])

        self._results = OrderedDict()

        if self.cfg.TEST.AQE.ENABLED:
            self.logger.info("Test with AQE setting")
            qe_time = self.cfg.TEST.AQE.QE_TIME
            qe_k = self.cfg.TEST.AQE.QE_K
            alpha = self.cfg.TEST.AQE.ALPHA
            query_features, gallery_features = aqe(query_features, gallery_features, qe_time, qe_k, alpha)

        if self.cfg.TEST.METRIC == "cosine":
            query_features = F.normalize(query_features, dim=1)
            gallery_features = F.normalize(gallery_features, dim=1)

        dist = self.cal_dist(self.cfg.TEST.METRIC, query_features, gallery_features)

        if p_att==None and q_att==None:
            raise False
        else:
            #type 2
            bias_dist = bias_dist / bias_dist.sum(-1).unsqueeze(-1)  # n x m x NUM_ATT
            if q_att == 0:
                bias_dist, _ = bias_dist[:,:,2:8].max(-1)  # n x m
            elif q_att == 1:
                bias_dist, _ = bias_dist[:, :, -7:].max(-1)  # n x m
            elif q_att == 2:
                bias_dist_1, _ = bias_dist[:, :, 2:8].max(-1)  # n x m
                bias_dist_2, _ = bias_dist[:, :, -7:].max(-1)  # n x m
                bias_dist = bias_dist_1 + bias_dist_2
            elif q_att == 3:
                bias_dist, _ = bias_dist.max(-1)  # n x m
            elif q_att == 4:

                gallery_fake_attributes = torch.cat(self.att_prec,dim=0)[self._num_query:]
                dist_fake_stack = query_real_attributes.unsqueeze(1) * gallery_fake_attributes.unsqueeze(
                    0)  # -1 means different, 1 means same

                dist_fake_stack = ((-1.0*dist_fake_stack)+1)/2.0  #-1 -> 1 ,1 -> 0,1 means different, 0 means same

                bias_dist = (bias_dist*dist_fake_stack).sum(-1)
            else:
                assert False

            bias_dist = (1.0+p_att*bias_dist)*dist

        if  self.cfg.TEST.RERANK.ENABLED:
            self.logger.info("Test with rerank setting")
            k1 = self.cfg.TEST.RERANK.K1
            k2 = self.cfg.TEST.RERANK.K2
            lambda_value = self.cfg.TEST.RERANK.LAMBDA
            q_q_dist = self.cal_dist(self.cfg.TEST.METRIC, query_features, query_features)
            g_g_dist = self.cal_dist(self.cfg.TEST.METRIC, gallery_features, gallery_features)
            re_dist = re_ranking(dist, q_q_dist, g_g_dist, k1, k2, lambda_value)
            re_dist = re_dist * bias_dist.numpy()
            query_features = query_features.numpy()
            gallery_features = gallery_features.numpy()
            cmc, all_AP, all_INP = evaluate_rank(re_dist, query_features, gallery_features,
                                                 query_pids, gallery_pids, query_camids,
                                                 gallery_camids, use_distmat=True)
        else:
            #dist = dist * bias_dist.numpy()
            dist = bias_dist.numpy()
            query_features = query_features.numpy()
            gallery_features = gallery_features.numpy()
            cmc, all_AP, all_INP = evaluate_rank(dist, query_features, gallery_features,
                                                 query_pids, gallery_pids, query_camids, gallery_camids,
                                                 use_distmat=True)
        mAP = np.mean(all_AP)
        mINP = np.mean(all_INP)
        for r in [1, 5, 10]:
            self._results['New_Rank-{}'.format(r)] = cmc[r - 1]
        self._results['New_mAP'] = mAP
        self._results['New_mINP'] = mINP

        if self.cfg.TEST.ROC_ENABLED:
            scores, labels = evaluate_roc(dist, query_features, gallery_features,
                                          query_pids, gallery_pids, query_camids, gallery_camids)
            fprs, tprs, thres = metrics.roc_curve(labels, scores)

            for fpr in [1e-4, 1e-3, 1e-2]:
                ind = np.argmin(np.abs(fprs - fpr))
                self._results["New_TPR@FPR={:.0e}".format(fpr)] = tprs[ind]

        return copy.deepcopy(self._results)

    def evaluate(self):
        if comm.get_world_size() > 1:
            raise False
        else:
            features = self.features
            pids = self.pids
            camids = self.camids

        features = torch.cat(features, dim=0)
        # query feature, person ids and camera ids
        query_features = features[:self._num_query]
        query_pids = np.asarray(pids[:self._num_query])
        query_camids = np.asarray(camids[:self._num_query])

        # gallery features, person ids and camera ids
        gallery_features = features[self._num_query:]
        gallery_pids = np.asarray(pids[self._num_query:])
        gallery_camids = np.asarray(camids[self._num_query:])

        self._results = OrderedDict()

        if self.cfg.TEST.AQE.ENABLED:
            self.logger.info("Test with AQE setting")
            qe_time = self.cfg.TEST.AQE.QE_TIME
            qe_k = self.cfg.TEST.AQE.QE_K
            alpha = self.cfg.TEST.AQE.ALPHA
            query_features, gallery_features = aqe(query_features, gallery_features, qe_time, qe_k, alpha)

        if self.cfg.TEST.METRIC == "cosine":
            query_features = F.normalize(query_features, dim=1)
            gallery_features = F.normalize(gallery_features, dim=1)

        dist = self.cal_dist(self.cfg.TEST.METRIC, query_features, gallery_features)

        if self.cfg.TEST.RERANK.ENABLED:
            self.logger.info("Test with rerank setting")
            k1 = self.cfg.TEST.RERANK.K1
            k2 = self.cfg.TEST.RERANK.K2
            lambda_value = self.cfg.TEST.RERANK.LAMBDA
            q_q_dist = self.cal_dist(self.cfg.TEST.METRIC, query_features, query_features)
            g_g_dist = self.cal_dist(self.cfg.TEST.METRIC, gallery_features, gallery_features)
            re_dist = re_ranking(dist, q_q_dist, g_g_dist, k1, k2, lambda_value)
            query_features = query_features.numpy()
            gallery_features = gallery_features.numpy()
            cmc, all_AP, all_INP = evaluate_rank(re_dist, query_features, gallery_features,
                                                 query_pids, gallery_pids, query_camids,
                                                 gallery_camids, use_distmat=True)
        else:
            query_features = query_features.numpy()
            gallery_features = gallery_features.numpy()
            cmc, all_AP, all_INP = evaluate_rank(dist, query_features, gallery_features,
                                                 query_pids, gallery_pids, query_camids, gallery_camids,
                                                 use_distmat=False)
        mAP = np.mean(all_AP)
        mINP = np.mean(all_INP)
        for r in [1, 5, 10]:
            self._results['Rank-{}'.format(r)] = cmc[r - 1]
        self._results['mAP'] = mAP
        self._results['mINP'] = mINP

        if self.cfg.TEST.ROC_ENABLED:
            scores, labels = evaluate_roc(dist, query_features, gallery_features,
                                          query_pids, gallery_pids, query_camids, gallery_camids)
            fprs, tprs, thres = metrics.roc_curve(labels, scores)

            for fpr in [1e-4, 1e-3, 1e-2]:
                ind = np.argmin(np.abs(fprs - fpr))
                self._results["TPR@FPR={:.0e}".format(fpr)] = tprs[ind]

        return copy.deepcopy(self._results)

def visiual_average_att(att_name, feature_mask, real_attributes,img_size=(384, 192), output_dir="./",positive=True):

    n, c, h, w = feature_mask.size()
    feature_mask = feature_mask.reshape(n, c, h * w)  # n x c x h*w
    if positive:
        feature_mask = feature_mask*((real_attributes+1.0)/2.0).unsqueeze(dim=-1)
    else:
        feature_mask = feature_mask * ((-real_attributes + 1.0) / 2.0).unsqueeze(dim=-1)
    # feature_max, _ = feature_mask.max(dim=-1)  # n x c
    # feature_max = feature_max.unsqueeze(-1)  # n x c  x 1
    # feature_min, _ = feature_mask.min(dim=-1)  # n x c
    # feature_min = feature_min.unsqueeze(-1)  # n x c  x 1
    #
    # feature_mask = (feature_mask - feature_min) / (feature_max - feature_min)  # n x c x h*w

    values, _ = feature_mask.sort(dim=-1, descending=True)
    max_v = values[:, :, h*w//3].unsqueeze(-1)
    feature_mask = torch.where( feature_mask > max_v, torch.ones_like(feature_mask), torch.zeros_like(feature_mask))

    feature_mask = feature_mask.sum(dim=0)  # c x h*w
    feature_max, _ = feature_mask.max(dim=-1)  # c
    feature_max = feature_max.unsqueeze(-1)  # c x 1
    feature_min, _ = feature_mask.min(dim=-1)  # c
    feature_min = feature_min.unsqueeze(-1)  # c x 1

    feature_mask = (feature_mask - feature_min) / (feature_max - feature_min).clamp(min=1e-6)  # c x h*w
    feature_mask = feature_mask.reshape(c, h, w).unsqueeze(0)  # 1 x c x h x w

    feature_mask = torch.nn.Upsample(size=img_size, scale_factor=None, mode='bilinear', align_corners=False)(
        feature_mask)
    feature_mask = feature_mask[0]
    len_att = len(att_name)
    num_col = 5 #int(len_att)

    num_raw = ((len_att - 1) // num_col + 1) * 2
    plt.figure(figsize=(num_col, num_raw), dpi=224)

    for i in range(0, num_raw, 2):
        for j in range(num_col):
            if i // 2 * num_col + j >= len_att:
                break
            plt.subplot(num_raw, num_col, i * num_col + j + 1)
            plt.axis('off')  # 去掉坐标轴
            plt.text(0.5, 0.5, att_name[i // 2 * num_col + j], fontsize=8, ha='center', va="center")

    for i in range(1, num_raw, 2):
        for j in range(num_col):
            if i // 2 * num_col + j >= len_att:
                break
            plt.subplot(num_raw, num_col, i * num_col + j + 1)
            plt.axis('off')  # 去掉坐标轴

            plt.imshow(feature_mask[i // 2 * num_col + j], cmap='nipy_spectral')
            #plt.imshow(cv2.applyColorMap(np.uint8(255 * feature_mask[i // 2 * num_col + j]), cv2.COLORMAP_JET))
            if j==num_col-1:
                plt.colorbar()

    plt.tight_layout()
    if positive:
        plt.savefig(os.path.join(output_dir, str(time.time())+"_positive_att_mask.jpg"))
    else:
        plt.savefig(os.path.join(output_dir, str(time.time())+"_negative_att_mask.jpg"))

    plt.clf()
    plt.close()

def visualization_savefig(path_names, att_names, imgs, att_mask, feature_distance, real_attribute, fake_attribute,output_dir="./",only_hot_map=False):
    """
    path_names:[str,...,str]
    att_names: [str,...,str]
    imgs: tensor -> n x 3 x 224 x 224
    att_mask: tensor -> n x 23 x 224 x 224
    feature_distance: tensor  -> (1+23+1) x 4
    real_attribute: tensor -> 23 x n
    fake_attribute: tensor -> 23 x n
    output_dir: default: "./"
    """
    #self.cfg.MODEL.HEADS.IN_FEAT
    # b_att,c_att,h_att,w_att = att_mask.size()
    # att_mask_use = torch.zeros(size=(b_att,c_att//IN_FEAT,h_att,w_att))
    # for i in range(c_att//IN_FEAT):
    #     att_mask_use[:,i,:,:] = att_mask[:,i*IN_FEAT:(i+1)*IN_FEAT,:,:].mean(dim=1)
    # att_mask = att_mask_use
    torch.save({'path_names':path_names,'imgs':imgs,'att_mask':att_mask,'real_attribute':real_attribute}, os.path.join(output_dir, path_names[0].replace(".", "_") + "mask.pth"))

    _,_,h, w = imgs.size()
    att_mask = torch.nn.Upsample(size=(h, w), scale_factor=None, mode='bilinear', align_corners=False)(att_mask)
    # rank of feature
    feature_distance_for_rank = feature_distance[1:-1]
    _, indices = feature_distance_for_rank.sort(dim=0, descending=True)
    _, indices = indices.sort(dim=0)
    indices = indices + 1

    num_low = len(att_names) + 2  # 25
    num_col = len(path_names) * 2  # 10
    plt.figure(figsize=(num_col, num_low), dpi=224)

    # name of pic
    for i, j in zip([0] * (num_col // 2), list(range(0, num_col, 2))):
        plt.subplot(num_low, num_col, i * num_col + j + 2)
        plt.axis('off')  # 去掉坐标轴
        plt.text(0.5, 0.5, path_names[j // 2], fontsize=6, ha='center', va="center")

    # name of title
    plt.subplot(num_low, num_col, num_col + 1)
    plt.axis('off')  # 去掉坐标轴
    str_show = "orginal"
    plt.text(0.5, 0.5, str_show, fontsize=14, ha='center', va="center")

    # name of att
    for i, j in zip(list(range(2, num_low, 1)), [0] * (num_low - 2)):
        plt.subplot(num_low, num_col, i * num_col + j + 1)
        plt.axis('off')  # 去掉坐标轴
        str_show = att_names[i - 2]
        str_show += "\n real_att: " + str(int(real_attribute[i - 2, 0]))
        str_show += "\n fake_att: " + str(float(fake_attribute[i - 2, 0]))[:6]
        plt.text(0.5, 0.5, str_show, fontsize=6, ha='center', va="center")

    # image
    for i in list(range(1, num_low, 1)):
        for j in list(range(0, num_col, 2)):
            plt.subplot(num_low, num_col, i * num_col + j + 2)
            plt.axis('off')  # 去掉坐标轴
            if i == 1:
                plt.imshow(imgs[j // 2].permute(1, 2, 0))
            else:
                min_att_mask = att_mask[j // 2, i - 2].min()
                max_att_mask = att_mask[j // 2, i - 2].max()
                att_mask_use = (att_mask[j // 2, i - 2] - min_att_mask) / (max_att_mask - min_att_mask)  # 0~1

                if only_hot_map:
                    # For a better viewing experience
                    plt.imshow(att_mask_use.pow(2), cmap='nipy_spectral')
                    #plt.imshow(att_mask_use, cmap='nipy_spectral')
                else:

                    plt.imshow((imgs[(j // 2)]).permute(1, 2, 0))
                    plt.imshow(att_mask_use.pow(2), cmap='nipy_spectral',alpha=0.4)
                    #plt.imshow(att_mask_use, cmap='nipy_spectral',alpha=0.4)

    # distance
    for j in list(range(2, num_col, 2)):
        # all_gap = 0
        for i in list(range(1, num_low, 1)):

            plt.subplot(num_low, num_col, i * num_col + j + 1)
            plt.axis('off')  # 去掉坐标轴
            # print(i-1,j//2-1)
            if i == 1:
                #TODO CXD ABS .abs()
                gap_fake_real = (feature_distance[-1, j // 2 - 1] - feature_distance[i - 1, j // 2 - 1]) / \
                                feature_distance[i - 1, j // 2 - 1] * 100
                str_show = str(float(feature_distance[i - 1, j // 2 - 1]))[:8] + "\n\n gap: " + str(
                    float(gap_fake_real))[:6] + "%"
                str_show += "\n var:"+str(float((feature_distance[1:-1, j // 2 - 1]/feature_distance[-1:, j // 2 - 1]).var()))[:8]

                str_show += "\n real_num:"+str(float(((1-(real_attribute[:,j // 2] * real_attribute[:,0]))/2).sum()))[:6]
                plt.text(0.5, 0.5, str_show, fontsize=6, ha='center', va="center")
            else:

                outstand_gap = feature_distance[i - 1, j // 2 - 1] / feature_distance[-1, j // 2 - 1] * 100
                rank_att = str(float(indices[i - 2, j // 2 - 1]))
                # all_gap+=outstand_gap
                str_show = str(float(feature_distance[i - 1, j // 2 - 1]))[:8] + "\n" + str(float(outstand_gap))[
                                                                                :6] + "%\n rank: " + rank_att
                str_show += "\n real_att: " + str(int(real_attribute[i - 2, j // 2]))
                #str_show += "\n fake_att: " + str(float(fake_attribute[i - 2, j // 2]))[:6]
                plt.text(0.5, 0.5, str_show, fontsize=6, ha='center', va="center")

        # print(all_gap)
    if only_hot_map:
        plt.savefig(os.path.join(output_dir, path_names[0].replace(".", "_") + "only_hot_map.jpg"))
    else:
        plt.savefig(os.path.join(output_dir, path_names[0].replace(".", "_") + "notonly_hot_map.jpg"))
    torch.save(feature_distance,os.path.join(output_dir, path_names[0].replace(".", "_") + ".pth"))
    plt.clf()
    plt.close()


'''

Sonic RGB LED for "Dr Who" themed sound & light box

Based on a very cut down version of :

    # http://www.lightshowpi.com/
    #
    # Author: Todd Giles (todd@lightshowpi.com)
    # Author: Chris Usey (chris.usey@gmail.com)
    # Author: Ryan Jennings
    # Jan 12 2014 - latest additions / mutilations by Scott Driscoll
    # CuriousInventor.com | https://plus.google.com/+ScottDriscoll
    #
    # Licensed under the BSD license.  See full license in LICENSE file.

'''

from __future__ import print_function
import RPi.GPIO as GPIO
import time
from multiprocessing import Process
import wave
import decoder
import numpy as np
import alsaaudio as audio
import os
import threading

# Number of frequency channels
GPIOLEN = 3

# How much data to read at once (must be 2**n)
CHUNK_SIZE = 2048

GPIO.setmode(GPIO.BOARD)
GPIO.setwarnings(False)

class SonicRGB(object):

    number = 0

    def __init__(self, red, green, blue, commonCathode=False,
                 pwmFrequency=100, cutoffs=None):
        self.red = red
        self.green = green
        self.blue = blue
        self.commonCathode = commonCathode
        self.pwmFrequency = pwmFrequency
        self.cutoffs = cutoffs
        if commonCathode:
            self.ON = 0
        else:
            self.ON = 100
        self.OFF = 100 - self.ON

        self.playing = None

        # Use numbering based on P1 header
        GPIO.setup(self.red, GPIO.OUT)
        GPIO.setup(self.green, GPIO.OUT)
        GPIO.setup(self.blue, GPIO.OUT)

        GPIO.output(self.red, GPIO.HIGH)
        GPIO.output(self.green, GPIO.HIGH)
        GPIO.output(self.blue, GPIO.HIGH)

        # PWM
        self.pwm = []
        self.pwm.append(GPIO.PWM(self.red, self.pwmFrequency))
        self.pwm.append(GPIO.PWM(self.green, self.pwmFrequency))
        self.pwm.append(GPIO.PWM(self.blue, self.pwmFrequency))
        for p in self.pwm:
            p.start(self.OFF)

        # id of current player
        self.curId = None

        # Flag for restaring current track
        self._restart = False

        # Flag to terminate current track
        self.stopNow = False

    def stop(self):
        if self.playing:
            self.stopNow = True
            self.thread.join()

    def play(self, id, track, finished):

        # If we're already playing a track, interrupt it and wait
        # for everything to be cleared
        if self.playing:
            self._restart = True
            return

        # id of caller
        self.curId = id

        # Set up audio
        if track.lower().endswith('.wav'):
            self.musicFile = wave.open(track, 'r')
        else:
            self.musicFile = decoder.open(track)

        self.sample_rate = self.musicFile.getframerate()
        self.num_channels = self.musicFile.getnchannels()

        self.output = audio.PCM(audio.PCM_PLAYBACK, audio.PCM_NORMAL)
        self.output.setchannels(self.num_channels)
        self.output.setrate(self.sample_rate)
        self.output.setformat(audio.PCM_FORMAT_S16_LE)
        self.output.setperiodsize(CHUNK_SIZE)

        self.frequency_limits = self._calculate_channel_frequency(50, 15000)

        # Start playing...
        self.playing = True
        self.thread = threading.Thread(target=self._play, args=[finished])
        self.thread.start()

    def _play(self, finished):
        for p in self.pwm:
            p.ChangeDutyCycle(self.OFF)

        data = self.musicFile.readframes(CHUNK_SIZE)
        while data != '' and not self.stopNow:
            try:
                self.output.write(data)
            except:
                break

            values = self._calculate_levels(data, self.sample_rate,
                                            self.frequency_limits)

            if self.commonCathode:
                for p, val in zip(self.pwm, values):
                    p.ChangeDutyCycle(int(100 - val))
            else:
                for p, val in zip(self.pwm, values):
                    p.ChangeDutyCycle(int(val))

            if self._restart:
                try:
                    self.musicFile.rewind()
                except:
                    pass
                self._restart = False
            data = self.musicFile.readframes(CHUNK_SIZE)

        self.musicFile.close()
        for p in self.pwm:
            p.ChangeDutyCycle(self.OFF)
        finished()
        self.playing = False
        self.curId = None
        self.stopNow = False
        return 0

    def busy(self, id):
        if self.curId:
            return self.curId != id
        else:
            return False

    def restart(self):
        self._restart = True

    # Light up in red
    def redAlert(self):
        self.pwm[0].ChangeDutyCycle(self.ON)
        self.pwm[1].ChangeDutyCycle(self.OFF)
        self.pwm[2].ChangeDutyCycle(self.OFF)

    def _calculate_levels(self, data, sample_rate, frequency_limits):
        '''Calculate frequency response for each channel

        Initial FFT code inspired from the code posted here:
        http://www.raspberrypi.org/phpBB3/viewtopic.php?t=35838&p=454041

        Optimizations from work by Scott Driscoll:
        http://www.instructables.com/id/Raspberry-Pi-Spectrum-Analyzer-with-RGB-LED-Strip-/
        '''

        # create a numpy array. This won't work with a mono file, stereo only.
        data_stereo = np.frombuffer(data, dtype=np.int16)

        if (self.num_channels == 2):
            # data has two channels and 2 bytes per channel
            data = np.empty(len(data) / 4)
            # pull out the even values, just using left channel
            data[:] = data_stereo[::2]
        else:
            data = data_stereo

        # if you take an FFT of a chunk of audio, the edges will look like
        # super high frequency cutoffs. Applying a window tapers the edges
        # of each end of the chunk down to zero.
        window = np.hanning(len(data))
        data = data * window

        # Apply FFT - real data
        fourier = np.fft.rfft(data)

        # Remove last element in array to make it the same size as CHUNK_SIZE
        fourier = np.delete(fourier, len(fourier) - 1)

        # Calculate the power spectrum
        power = np.abs(fourier) ** 2

        result = [0 for i in range(GPIOLEN)]
        for i in range(GPIOLEN):
            # take the log10 of the resulting sum to approximate how human ears
            # perceive sound levels
            result[i] = np.sum(power[self._piff(frequency_limits[i][0], sample_rate):
                                     self._piff(frequency_limits[i][1], sample_rate):1])
        result = np.clip(result, 1.0, 1.0e20)
        mag = 0.01 * np.sqrt(np.dot(result, result))
        result = result / mag
        # print(result)
        #result = np.clip(100 * result / np.max(result), 0, 100)
        return result

    def _piff(self, val, sample_rate):
        '''Return the power array index corresponding to a particular frequency.'''
        return int(CHUNK_SIZE * val / sample_rate)

    def _calculate_channel_frequency(self, min_frequency, max_frequency):
        '''Calculate frequency values for each channel'''

        if self.cutoffs:
            f0, f1, f2, f3 = self.cutoffs
            return [[f0, f1], [f1, f2], [f2, f3]]

        # How many channels do we need to calculate the frequency for
        channel_length = GPIOLEN

        octaves = (np.log(max_frequency / min_frequency)) / np.log(2)
        octaves_per_channel = octaves / channel_length
        frequency_limits = []
        frequency_store = []

        frequency_limits.append(min_frequency)
        for i in range(1, GPIOLEN + 1):
            frequency_limits.append(
                frequency_limits[-1] * 2 ** octaves_per_channel)
        for i in range(0, channel_length):
            frequency_store.append(
                (frequency_limits[i], frequency_limits[i + 1]))

        return frequency_store


"""Load an existing gameworld or begin anew"""

from copy import deepcopy
from random import choices, choice, randint, shuffle
from operator import add
import matplotlib.pyplot as plt
from numpy import full as npf
import dork.game_utils.factory_data as factory_data
# pylint: disable=protected-access

class ItemFactory:
    """Generates a random named item with randomized stats"""

    items = factory_data.ITEMS
    names = factory_data.NAMES
    sequence = factory_data.SEQUENCE

    types = items["types"]
    condition = items["condition"]
    material = items["material"]

    posessive = names["posessive"]
    nonposessive = names["nonposessive"]
    suffixes = names["suffixes"]
    abstract = names["abstract"]
    adjectives = names["adjectives"]

    @staticmethod
    def build(weights=None):
        """generate a random item

        Creates an item name, type, and location into game world

        Args:
            item_name (str): name of randomly generated item
            item_type (str): tpye of randomly generated item

        returns:
            item_name (str): returns the name of the randomly generated item
            item_type (str): returns the type of the randomly generated item

        """

        weights = {
            "player": [8, 0, 0, 7, 5, 10]
        }.get(weights, [8, 35, 3, 7, 5, 10])

        item_type = choice(choices(
            population=list(ItemFactory.types.keys()),
            weights=weights,
            k=len(list(ItemFactory.types.keys()))
        ))

        item_name = choice(choices(
            population=ItemFactory.types[item_type],
            k=len(ItemFactory.types[item_type])
        ))

        return ItemFactory._forge(item_name, item_type)

    @staticmethod
    def _generate(stats, item_name, item_type):
        return {
            "name": item_name,
            "type": item_type,
            "description": "",
            "stats": stats
        }

    @staticmethod
    def _stats(item_name, item_type):
        stats = factory_data.stats(item_type.split()[0])
        return ItemFactory._generate(stats, item_name, item_type)

    @staticmethod
    def _forge(item_name, item_type):
        new_name = []
        build = ItemFactory.sequence[item_type]

        seq = choice(choices(
            population=build["seq"],
            weights=build["w"],
            k=len(build["seq"])
        ))

        for lists in seq:
            if isinstance(lists, dict):
                this_list = lists.get(
                    item_type, lists.get("usable", ['']))
            elif lists:
                this_list = lists
            else:
                this_list = ['']

            this_word = choice(choices(
                population=this_list,
                k=len(this_list)
            ))

            if this_word:
                if this_word in ItemFactory.suffixes:
                    new_name[-1] += this_word
                    item_type = f"legendary {item_name}"
                else:
                    new_name.append(this_word)
            else:
                new_name.append(item_name)

        item_name = " ".join(new_name)
        return ItemFactory._stats(item_name, item_type)

class PlayerFactory:
    """Generate players for a room"""

    @staticmethod
    def build(i, room):
        """Make a player, give them items

        Extended Description

        Args:
            arg1 (str): description

        returns:
            arg1 (str): description

        """
        firsts = factory_data.FIRST_NAMES
        rand_first = choice(firsts)
        lasts = factory_data.LAST_NAMES
        rand_last = choice(lasts)
        new_player = {
            "name": rand_first + " " + rand_last,
            "description": f"player {i} description",
            "location": room["name"],
            "inventory": {},
            "equipped": []
        }

        for _ in range(randint(1, 3)):
            new_item = ItemFactory.build("player")
            new_player["inventory"][new_item["name"]] = new_item

        for key, val in new_player["inventory"].items():
            if val["stats"]["equipable"]:
                new_player["equipped"].append(key)

        return new_player

class RoomFactory:
    """Generate rooms for a given maze"""

    #  N, S and E, W are backwards because numpy uses column-order
    moves = {
        "north": (1, 0), "south": (-1, 0),
        "east": (0, 1), "west": (0, -1),
    }

    @staticmethod
    def build(maze, rooms):
        """build a room

        Creates an instance of a room

        Args:
            _make_rooms (): creates an instance of a room

        returns:
            _make_rooms (): returns an instance of a room

        """

        RoomFactory.maze = maze
        RoomFactory.rooms = rooms
        RoomFactory.worldmap = {}
        return RoomFactory._make_rooms()

    @staticmethod
    def _make_rooms():

        list_of_keys = factory_data.ROOMS
        shuffle(list_of_keys)
        list_of_adjtvs = factory_data.NAMES["adjectives"]
        shuffle(list_of_adjtvs)
        list_of_abstract = factory_data.NAMES["abstract"]
        shuffle(list_of_abstract)

        i = 0

        for room in RoomFactory.rooms:
            if i == 0:
                x, y = room
                new_room = {
                    "number": f"room 0",
                    "name": f"Entrance",
                    "description": factory_data.DEFAULT_ROOMS["Entrance"],
                    "coordinates": [x, y],
                    "adjacent": {},
                    "players": {},
                    "inventory": {},
                }
            elif i < len(RoomFactory.rooms) - 1:
                rand = list_of_keys[i]
                x, y = room
                new_room = {
                    "number": f"room {i}",
                    "name": list_of_adjtvs[i] + rand + list_of_abstract[i],
                    "description": "The " + list_of_adjtvs[i] + rand +
                                   list_of_abstract[i],
                    "coordinates": [x, y],
                    "adjacent": {},
                    "players": {},
                    "inventory": {},
                }
            else:
                x, y = room
                new_room = {
                    "number": f"room "+str(len(RoomFactory.rooms)),
                    "name": f"End",
                    "description": factory_data.DEFAULT_ROOMS["End"],
                    "coordinates": [x, y],
                    "adjacent": {},
                    "players": {},
                    "inventory": {},
                }

            for _ in range(randint(1, 7)):
                new_item = ItemFactory.build()
                new_room["inventory"][new_item["name"]] = new_item

            for _ in range(randint(0, 2)):
                new_player = PlayerFactory.build(i, new_room)
                new_room["players"][new_player["name"]] = new_player

            RoomFactory.worldmap[room] = new_room
            i += 1

        return RoomFactory._get_adj()

    @staticmethod
    def _get_adj():
        for coord, room in RoomFactory.worldmap.items():
            for direction in RoomFactory.moves:
                searching = True
                position = coord
                while searching:
                    position = tuple(
                        map(add, position, RoomFactory.moves[direction]))
                    if RoomFactory.maze[position] == MazeFactory.wall_color:
                        room["adjacent"][direction] = None
                        searching = False
                    elif RoomFactory.maze[position] in \
                            [MazeFactory.room_color, MazeFactory.player_color]:
                        room["adjacent"][direction] = \
                            RoomFactory.worldmap[position]["number"]
                        searching = False

        for coord, room in deepcopy(RoomFactory.worldmap).items():
            new_room = RoomFactory.worldmap.pop(coord)
            RoomFactory.worldmap[new_room.pop("number")] = new_room

        return RoomFactory.worldmap

class MazeFactory:
    """Generate a maze with rooms on intersections, corners, and dead-ends"""

    wall_color, path_color, room_color, player_color = (-2, 2, 1, 0)
    moves = factory_data.MOVES
    rules = factory_data.rules(wall_color, path_color)

    @staticmethod
    def draw(maze):
        """display the maze"""

        plt.figure(figsize=(len(maze[0])//2, len(maze)//2))
        plt.pcolormesh(maze, cmap=plt.cm.get_cmap("tab20b"))
        plt.axis("equal")
        plt.axis("off")
        plt.ion()
        plt.show()

    @staticmethod
    def update(maze):
        """update the maze display"""

        plt.pcolormesh(maze, cmap=plt.cm.get_cmap("tab20b"))
        plt.axis("equal")
        plt.axis("off")
        plt.draw()

    # pylint: disable=R0914
    @staticmethod
    def build():
        """generate a maze"""

        x = choice([10, 12, 14, 18])
        y = 148//x

        maze = npf((x+1, y+1), MazeFactory.wall_color)
        grid = [(i, j) for i in range(1, x+1, 2) for j in range(1, y+1, 2)]
        path = [choice(grid)]
        rooms = []
        position = path[0]
        grid.remove(position)

        while grid:
            n = len(path)
            nsew = []
            for move in MazeFactory.moves:
                nsew.append([
                    tuple(map(add, move[0], position)),
                    tuple(map(add, move[1], position))
                ])
            shuffle(nsew)
            for probe in nsew:
                if probe[0] in grid:
                    maze[probe[0]] = MazeFactory.path_color
                    maze[probe[1]] = MazeFactory.path_color
                    grid.remove(probe[0])
                    path.extend(probe)
                    break
            if n == len(path):
                position = path[max(path.index(position)-1, 1)]
            else:
                position = path[-1]

        for coord in path:
            i, j = coord
            neighbors = [
                maze[i-1, j],
                maze[i+1, j],
                maze[i, j-1],
                maze[i, j+1]
            ]
            if neighbors in MazeFactory.rules:
                rooms.append(coord)
                maze[coord] = MazeFactory.room_color
            maze[rooms[0]] = MazeFactory.player_color

        return {
            "maze": maze.tolist(),
            "rooms": RoomFactory.build(maze, rooms)
        }