"""Markov Chain Monte Carlo for trDesign."""

# native
from datetime import datetime
import time

# lib
import numpy as np
from numpy.random import choice
import torch
from torch.cuda.amp import autocast
import math
import random
import time

# pkg
from losses import *  # pylint: disable=wildcard-import, unused-wildcard-import
from tr_Rosetta_model import trRosettaEnsemble, prep_seq
from utils import aa2idx, distogram_distribution_to_distogram, idx2aa, plot_progress
import config as cfg
from utils import definegroupbydist, plot_muts

def v(torch_value):
    """Return a detached value, if possible."""
    try:
        return torch_value.cpu().detach().item()
    except Exception:
        return torch_value

def motifsort(elem):
    return elem[4]

def placemotifs(motifs, seq_L, sequence, mode = 0):
    """Randomly position discontinous motifs and check if valid. motif = [start, end, length, restraints, group, newstart, newend]"""
    print("Placing motifs...")
    valid = False
    i = 0
    sum = 0

    for m in motifs:
        sum += m[2]
    if (sum > seq_L):
        print("Sequence too short")
        return placemotifs(motifs[:-1], seq_L, sequence, mode)
    elif sum == seq_L:
        pos = 0
        mode = -3
        for m in motifs:
            m[5] = pos
            m[6] = m[5] + m[2] - 1
            pos = pos + m[2]

    while not valid:
        if mode == 0 or mode == 1:
            for m in motifs[:]:
                m[5] = np.random.randint(0,seq_L-m[2]+1)
                m[6] = m[5]+m[2]-1
        elif mode == -2:
            print("predefined motifs")
            seq_con = ""
            for i in range(0,seq_L):
                restraint = "-"
                for m in motifs:
                    if i >= m[5] and i <= m[6]:
                        mi = i - m[5] #local index
                        c = m[3][mi]  #con type
                        if c in cfg.sequence_restraint_letters:
                            restraint = sequence[i]
                        continue

                seq_con = seq_con + restraint

            return motifs, seq_con
        elif mode == -3:
            break;
        elif mode == 2:
            spacing = (seq_L-sum)/(1+len(motifs))
            pos = int(spacing)
            for m in motifs[:]:
                    buffer = int(abs(np.random.normal(spacing,spacing)))
                    m[5] = pos
                    m[6] = m[5]+m[2]-1
                    pos = m[6] + buffer

        elif mode == 2.1:
            spacing = int((seq_L-sum-6)/(len(motifs)-1)) #calculate spacing for even placement with short unrestrained tails
            pos = int(seq_L-sum-((len(motifs)-1)*spacing)) #center the motifs
            for m in motifs[:]:
                    m[5] = pos
                    m[6] = m[5]+m[2]-1
                    pos = m[6]+spacing
        elif mode == 2.2:
            motifs[0][5] = 0
            motifs[0][6] = motifs[0][2]-1
            motifs[-1][6] = seq_L - 1
            motifs[-1][5] = seq_L - motifs[-1][2]
            last = motifs[-1]
            spacing = int((seq_L-sum)/(len(motifs)-1))
            pos = motifs[0][6]
            for m in motifs[1:-1]:
                buffer = int(abs(np.random.normal(spacing,spacing)))
                pos += 1 + buffer
                m[5] = pos
                m[6] = m[5] + m[2] - 1
                pos = m[6]

        elif mode == 3:
            motifs.sort(key=motifsort)
            for m in motifs[:]: #set all same group
                m[4] = 0
            return placemotifs(motifs, seq_L, sequence, mode = 2)
        elif mode == 4 or mode == 5:
            motifs = definegroupbydist(motifs,cfg.target_motif_path, mode)
            return placemotifs(motifs, seq_L, sequence, mode = 3)
        else: return placemotifs(motifs, seq_L, sequence, mode = 1)

        #check if motifs are valid
        valid = True
        i = i + 1
        for m1 in motifs[:]:
            if m1[6] > seq_L - 1:
                print("ends outside")
                valid = False
                continue
            for m2 in motifs[:]:
                if m1 == m2: continue
                if (m1[5] >= m2[5]) and (m1[5] <= m2[6]):
                    valid = False
                    break
                elif m1[6] >= m2[5] and m1[6] <= m2[6]:
                    valid = False
                    break

        if i > 20000: #if not valid, place motifs manually
            print("No valid motif placements found, attempting sequential positions")
            random.shuffle(motifs)
            rest = seq_L - sum
            buffer = math.floor(rest/(len(motifs)+1))
            pos = buffer
            for m in motifs:
                m[5] = pos
                m[6] = m[5] + m[2]-1
                pos += m[2] + buffer
            break

    seq_con = ""

    constrain_seq = False
    if sequence is not None:
        for i in range(0,seq_L):
            restraint = "-"
            for m in motifs:
                if i >= m[5] and i <= m[6]:
                    mi = i - m[5] #local index
                    c = m[3][mi]  #con type
                    if c in cfg.sequence_restraint_letters:
                        si = m[0]+mi  #template index
                        restraint = sequence[si]
                        constrain_seq = True
                    continue

            if i == 0 and cfg.first_residue_met:
                restraint = "M"

            seq_con = seq_con + restraint

    if not constrain_seq:
        seq_con = None

    return motifs, seq_con

def createmask(motifs, seq_L, save_dir, is_site_mask = False):
    """Create mask for discontinous motifs"""

    #setup mask size
    motif_mask_g = np.zeros((seq_L, seq_L))
    motif_mask = np.zeros((seq_L, seq_L))

    for m1 in motifs[:]:
        for i in range(m1[5],m1[6]+1):
            c1 = m1[3][i-m1[5]]
            if c1 in cfg.structure_restraint_letters:  #contraint is structural?
                for m2 in motifs[:]:
                    if math.fabs(m2[4]-m1[4]) > 1: #motifs are not in restrained groups?
                        continue
                    for j in range(m2[5],m2[6]+1):
                        c2 = m2[3][j-m2[5]]
                        if c2 in cfg.structure_restraint_letters:
                            v1 = cfg.structure_restraint_mask_values[c1]
                            v2 = cfg.structure_restraint_mask_values[c2]
                            value = np.clip(v1*v2, 1, 100)
                            motif_mask[i,j] = value
                            motif_mask_g[i,j] = value**0.5

    if not is_site_mask:
        plot_values = motif_mask_g.copy()
        plot_distogram(
            plot_values,
            save_dir / f"motif_mask_groups.jpg",
        )

    return motif_mask

class MCMC_Optimizer(torch.nn.Module):
    """Markov Chain Monte Carlo optimizer."""

    # We don't define `.forward()`, but it's nn-based.
    # pylint: disable=too-many-instance-attributes, abstract-method

    DEFAULT_ROOT_PATH = Path(__file__).parent.parent
    DEFAULT_MODEL_PATH = DEFAULT_ROOT_PATH / "models" / "trRosetta_models"
    DEFAULT_BGND_PATH = DEFAULT_ROOT_PATH / "backgrounds"
    DEFAULT_RESULTS_PATH = DEFAULT_ROOT_PATH / "results"

    def __init__(
        self,
        L,
        aa_weight,
        MCMC,
        native_frequencies,
        experiment_name,
        aa_valid,
        max_aa_index=20,
        sequence_constraint=None,
        target_motif_path=None,
        trRosetta_model_dir="models/trRosetta_models",
        background_distribution_dir="backgrounds",
        motifs=None,
        motifmode = 0,
        motif_weight = 1,
        bkg_weight = 1
    ):
        """Construct the optimizer."""
        super().__init__()
        self.results_dir = self.setup_results_dir(experiment_name)
        self.bkg_dir = background_distribution_dir
        self.structure_models = trRosettaEnsemble(
            trRosetta_model_dir
        )  # .share_memory()
        print(
            f"{self.structure_models.n_models} structure prediction models loaded to {d()}"
        )

        # General params:
        self.eps = 1e-7
        self.seq_L = L
        self.motifs = None
        self.motifmode = motifmode
        self.use_sites = cfg.use_sites
        print("Sequence Length: " + str(self.seq_L))
        if motifs is not None:
            _motifs,_seq_con = placemotifs(motifs, self.seq_L, sequence_constraint, mode=self.motifmode)
            self.motifs = _motifs
            if _seq_con is not None: sequence_constraint = _seq_con

        for m in self.motifs: print(m)
        print(sequence_constraint)

        # Setup MCMC params:
        self.beta, self.N, self.coef, self.M, self.MaxM = (
            MCMC["BETA_START"],
            MCMC["N_STEPS"],
            MCMC["COEF"],
            MCMC["M"],
            MCMC["MAX"],
        )
        self.aa_weight = aa_weight

        # Setup sequence constraints:
        self.aa_valid = aa_valid
        self.native_frequencies = native_frequencies

        self.seq_constraint = sequence_constraint
        if self.seq_constraint is not None:
            assert len(self.seq_constraint) == self.seq_L, \
            "Constraint length (%d) must == Seq_L (%d)" %(len(self.seq_constraint), self.seq_L)

            self.seq_constraint = (
                aa2idx(self.seq_constraint).copy().reshape([1, self.seq_L])
            )
            self.seq_constraint_indices = np.where(
                self.seq_constraint != max_aa_index, 1, 0
            )

        self.target_motif_path = target_motif_path
        self.setup_losses()

        # stats
        self.bad_accepts = []
        self.n_accepted_mutations = 0
        self.n_accepted_bad_mutations = 0
        self.best_metrics = {}
        self.best_step = 0
        self.best_sequence = None
        self.best_E = None
        self.step = 0
        self.motif_weight = motif_weight
        self.bkg_weight = bkg_weight

        self.matrix_setup()

    def matrix_setup(self):
        self.substitution_matrix = {}
        with open("src/" + cfg.MATRIXFILE) as reader:
            columns = reader.readline().strip().split()
            lines = reader.readlines()
            for aa in columns:
                if aa in cfg.RM_AA: continue
                if aa not in cfg.ALPHABET_core_str: continue
                self.substitution_matrix[aa] = {}

            for line in lines:
                data = line.strip().split()
                aa1 = data[0]
                if aa1 not in cfg.ALPHABET_core_str: continue
                sub = {}
                for i in range(0,len(columns)):
                    aa2 = columns[i]
                    if aa2 in cfg.RM_AA: continue
                    if aa2 not in cfg.ALPHABET_core_str: continue
                    value = data[i+1]
                    try: sub[aa2] = float(value)
                    except: sub[aa2] = 0

                self.substitution_matrix[aa1] = dict(sorted(sub.items(), key=lambda x: x[1]))

        if True: #convert to non negative values
            min_value = 1

            for a in self.substitution_matrix.items():
                for b in a[1].items():
                    if b[1] < min_value: min_value = b[1]

            if min_value < 0:
                offset = -min_value + 1

                for a in self.substitution_matrix.items():
                    for b in a[1].items():
                        self.substitution_matrix[a[0]][b[0]] += offset

                #for a in self.substitution_matrix.items():
                #    for b in a[1].items():
                #        self.substitution_matrix[a[0]][b[0]] = self.substitution_matrix[a[0]][b[0]]*self.substitution_matrix[a[0]][b[0]]

    def setup_results_dir(self, experiment_name):
        """Create the directories for the results."""
        results_dir = (
            self.DEFAULT_RESULTS_PATH
            / experiment_name
            / datetime.now().strftime("%Y-%m-%d_%H%M%S")
        )
        results_dir.mkdir(parents=True, exist_ok=True)
        (results_dir / "distogram_evolution").mkdir(parents=True, exist_ok=True)
        print(f"Writing results to {results_dir}")
        return results_dir

    def setup_losses(self):
        """Prepare the loss functions."""

        # Initialize protein background distributions:
        if cfg.BACKGROUND: self.bkg_loss = Structural_Background_Loss(self.seq_L, self.bkg_dir)
        self.aa_bkgr_distribution = torch.from_numpy(self.native_frequencies).to(d())

        # Motif-Loss:
        self.motif_weight = 1.00 #Multiplier for the loss contribution
        self.motif_mask = np.zeros((self.seq_L, self.seq_L)) #LxL tensor, int64, L currently specified in config.py. Cotains only zeros
        self.motif_mask = torch.from_numpy(self.motif_mask).long().to(d())

        #If motif target is specified. Motif target file is .npz LxL. L is overall protein length
        #Overrides previous motif_mask
        #New mask will be ones with a diagonal zero line
        if self.target_motif_path is not None:
            self.motif_mask = np.ones((self.seq_L, self.seq_L))

            if self.motifs is not None:
                self.motif_mask = createmask(self.motifs, self.seq_L, self.results_dir)

            self.motif_mask = torch.from_numpy(self.motif_mask).long().to(d())
            self.motif_mask.fill_diagonal_(0)
            self.motif_sat_loss = Motif_Satisfaction(
                self.target_motif_path, mask=self.motif_mask, save_dir=self.results_dir, motifs = self.motifs
            )

            self.site_weight = 0
            if self.use_sites:
                self.site_mask = createmask(self.motifs, self.seq_L, self.results_dir, is_site_mask = True)
                self.site_mask = torch.from_numpy(self.site_mask).long().to(d())
                self.site_mask.fill_diagonal_(0)

                self.site_sat_loss = Site_Satisfaction(
                    self.target_motif_path, mask=self.site_mask, motifs = self.motifs
                )

        # Apply the background KL-loss only under the hallucination_mask == 1 region
        self.hallucination_mask = 1 - torch.clip(self.motif_mask, 0, 1)
        self.hallucination_mask.fill_diagonal_(0)

    def loss(self, sequence, structure_predictions, msa1hot, track=False):
        """Compute the loss function."""

        # Top-prob: extremily expensive operation
        #TM_score_proxy = top_prob(structure_predictions['dist'], verbose=False)
        #TM_score_proxy = TM_score_proxy[0]  # We're running with batch_size = 1

        # Background KL-loss:
        if cfg.BACKGROUND:
            background_loss = self.bkg_loss(
                structure_predictions, hallucination_mask=self.hallucination_mask
            )
        else:
            background_loss = torch.tensor(0)

        # aa composition loss
        aa_samp = (
            msa1hot[0, :, :20].sum(axis=0) / self.seq_L + self.eps #take entire first sequence, all 1hot chars, and sum across length?
        )  # Get relative frequency for each AA
        aa_samp = (
            aa_samp / aa_samp.sum()
        )  # Normalize to turn into distributions (possibly redundant)
        loss_aa = (
            aa_samp
            * torch.log(aa_samp / (self.aa_bkgr_distribution + self.eps) + self.eps)
        ).sum()

        # Motif Loss:
        if self.target_motif_path is not None: #check if target motif has been specified
            motif_loss, motif_loss_pos = self.motif_sat_loss(structure_predictions) #motif_sat_loss = MotifSatisfaction object (nn object)
        else:
            motif_loss = 0 #no target motif = no loss

        if self.use_sites:
            site_loss = self.site_sat_loss(structure_predictions)
        else:
            site_loss = torch.tensor(0)

        # total loss
        loss_v = (
            self.bkg_weight * background_loss + self.aa_weight * loss_aa + self.motif_weight * motif_loss + self.site_weight * site_loss
        )

        metrics = {}
        if track:
            metrics["aa_weight"] = self.aa_weight
            metrics["background_loss"] = background_loss
            metrics["motif_loss"] = motif_loss
            metrics["site_loss"] = site_loss
            metrics["total_loss"] = loss_v
            #metrics["TM_score_proxy"] = TM_score_proxy

        return loss_v, metrics

    def metropolis(self, seq, seq_curr, E_curr, E):
        """Compute the Metropolis criterion."""
        accepted = True
        deltaE = E_curr - E

        # Metropolis criterion
        if E_curr < E:  # Lower energy, replace!
            seq = np.copy(seq_curr)
            E = E_curr
        else:  # Higher energy, maybe replace..
            if torch.exp((E - E_curr) * self.beta) > np.random.uniform():
                seq = np.copy(seq_curr)
                E = E_curr
                self.bad_accepts.append(1)
                self.n_accepted_bad_mutations += 1
            else:
                accepted = False
                self.bad_accepts.append(0)

        self.register_mutation_fitness(accepted, deltaE)

        # Update the best sequence:
        if E_curr < self.best_E:
            self.best_E = E_curr
            self.best_sequence = idx2aa(seq_curr[0])
            self.best_step = self.step

        return seq, E

    def mutate(self, seq):
        """Return a mutated version of the sequence."""
        seq_curr = np.copy(seq)

        # Introduce a random mutation using the allowed aa_types:
        if cfg.GRADIENT: idx = random.choices(range(self.seq_L), self.gradient, k=1)[0]
        else: idx = np.random.randint(self.seq_L)

        from_aa = idx2aa(seq_curr[0])[idx]

        #Perform mutation
        mutation_list = self.mutation_score[idx]

        if cfg.MATRIX:
            if len(mutation_list) == 0: mutation_list.append([from_aa, from_aa, 0, True, 0]) #Adds accepted dummy mutation

            #mut = self.mutation_score[idx]
            #mut_score = mut[2]
            #accepted = mut[3]
            #mins = min(self.mutation_score, key=lambda x: x[2])[2]
            #maxs = max(self.mutation_score, key=lambda x: x[2])[2]

            ################################
            ### SUBSTITUTION MATRIX MODE ###
            if cfg.MATRIX_MODE == 'probability':
                last_accepted = None
                verybad = ""

                for m in reversed(mutation_list):
                    if m[2] > 0.1:
                        verybad += m[1]
                    if m[3]:
                        last_accepted = m
                        break

                if last_accepted[2] < 0.005: p = 1
                else: p = 0

                for i in range(10):
                    list_of_candidates = [k for k,v in self.substitution_matrix[last_accepted[p]].items()]
                    probability_distribution = [v for k,v in self.substitution_matrix[last_accepted[p]].items()]
                    mutation = random.choices(list_of_candidates, probability_distribution, k=1)[0]
                    if mutation not in verybad: break

                seq_curr[0, idx] = aa2idx(mutation)

            #############################################
            ### AMINO ACID PROPERTY GROUP MATRIX MODE ###
            elif cfg.MATRIX_MODE == 'groups':
                recent_mutations = []

                for m in reversed(mutation_list):
                    if self.step - m[4] > self.seq_L * 10: break
                    recent_mutations.append(m)

                if len(recent_mutations) == 0:
                    seq_curr[0, idx] = np.random.choice(self.aa_valid)
                else:
                    group_scores = [[0,0,None]] * len(cfg.AA_PROPERTY_GROUPS)
                    i = 0
                    for k,v in cfg.AA_PROPERTY_GROUPS.items():
                        group_scores[i][2] = v
                        for m in recent_mutations:
                            if m[1] in v:
                                group_scores[i][1] += m[2]
                                group_scores[i][0] += 1
                        i += 1

                    best = ""
                    best_s = 1000

                    for i in range(len(group_scores)):
                        g = group_scores[i]
                        avg = g[1]/(g[0]+1)
                        if avg < best_s:
                            best = g[2]
                            best_s = avg
                        elif avg == best_s:
                            best += g[2]

                    options = []
                    for aa in best:
                        if aa in cfg.RM_AA: continue
                        options.append(aa)

                    seq_curr[0, idx] = aa2idx(np.random.choice(options))

            ################################
            ### MSA MUTATION MATRIX MODE ###
            elif cfg.MATRIX_MODE == 'msa':
                group = cfg.TEMPLATE_AA_PROPERTIES[idx]
                seq_curr[0, idx] = aa2idx(np.random.choice(group))

            ###########################
            ### MODE NOT RECOGNIZED ###
            else: seq_curr[0, idx] = np.random.choice(self.aa_valid)

        #######################
        ### NON MATRIX MODE ###
        else: seq_curr[0, idx] = np.random.choice(self.aa_valid)

        to_aa = idx2aa(seq_curr[0])[idx]

        if self.seq_constraint is not None:  # Fix the constraint:
            seq_curr = np.where(
                self.seq_constraint_indices, self.seq_constraint, seq_curr
            )

        if np.equal(seq_curr, seq).all(): # If the mutation did not change anything, retry
            self.reverse_log = False
            return self.mutate(seq)

        # Store mutation information
        self.current_mutations.append([idx, from_aa, to_aa])

        return seq_curr

    def select_mutation_options(self, matrix, threshold = 0, above = True):
        if above:
            return [k for k,v in matrix.items() if v >= threshold]
        else:
            return [k for k,v in matrix.items() if v <= threshold]

    def diff_to_weight(self, diff, increase = True):
        if increase:
            if diff == 0: return 1.25
            else: return 1.5 - 0.025*diff
        else: return 0.8 + 0.02 * diff

    def register_mutation_fitness(self, accepted, deltaE):
        if deltaE < -1: deltaE = 0
        else: deltaE = float(deltaE.cpu().detach().numpy())

        good = deltaE < 0

        for mut in self.current_mutations:
            self.good_mutations[mut[0]] += 1

            if good:
                self.good_accepts.append(1)
                for i in range(mut[0]-20,mut[0]+21):
                    if i >= 0 and i < self.seq_L:
                        self.gradient[i] *= self.diff_to_weight(abs(mut[0]-i))
            else:
                self.good_accepts.append(0)
                for i in range(mut[0]-10,mut[0]+11):
                    if i >= 0 and i < self.seq_L:
                        self.gradient[i] *= self.diff_to_weight(abs(mut[0]-i), increase = False)

            self.mutation_log.append([mut[0], mut[1], mut[2], self.substitution_matrix[mut[1]][mut[2]], deltaE, accepted])

            if accepted: self.n_accepted_mutations += 1

            self.mutation_score[mut[0]].append([mut[1], mut[2], deltaE, accepted, self.step]) #from, to, result, accepted, when

            if self.reverse_log:
                print("result:  " + str(mut) + " " + str(deltaE) + " | " + str(accepted))

        self.gradient = np.clip([i+0.01*(1-i) for i in self.gradient], 0.2, 20)
        self.current_mutations = []
        self.reverse_log = False

    def fixup_MCMC(self, seq):
        """Dynamically adjust the metropolis beta parameter to improve performance."""
        if self.step - self.best_step > min(self.N // 4,1000):
            # No improvement for a long time, reload the best_sequence and decrease beta:
            print("reload best seq")
            self.best_step = self.step
            self.beta = self.beta / (self.coef ** 2)
            seq = torch.from_numpy(
                aa2idx(self.best_sequence).reshape([1, self.seq_L])
            ).long()
            self.mutation_score = [[]] * self.seq_L

        elif np.mean(self.bad_accepts[-100:]) < 0.02:
            # There has been some progress recently, but we're no longer accepting any bad mutations...
            self.beta = self.beta / self.coef
        else:
            self.beta = self.beta * self.coef

        self.beta = np.clip(self.beta, 5, self.MaxM)

        #FIX DYNAMIC MATRIX
        if cfg.MATRIX_DYNAMIC: cfg.MATRIX_MODE = random.choice(['msa', 'probability', 'groups', 'none'])

        return seq

    @torch.no_grad()
    def run(self, start_seq):
        """Run the MCMC loop."""
        # pylint: disable=too-many-locals

        start_time = time.time()
        muttime = 0
        nntime = 0
        losstime = 0
        scoretime = 0
        misctime = 0

        # initialize with given input sequence
        print("Initial seq: ", start_seq)
        seq = aa2idx(start_seq).copy().reshape([1, self.seq_L])

        nsave = min(max(1, self.N // 20),50)
        E, E_tracker = np.inf, []
        self.bad_accepts = []
        self.good_accepts = []
        self.n_accepted_mutations = 0
        self.n_accepted_bad_mutations = 0
        self.best_metrics = {}
        self.best_step = 0
        self.best_sequence = start_seq
        self.best_E = E
        self.best_bkg = 0
        self.best_mtf = 0
        self.best_distogram_distribution = None
        self.gradient = [1] * self.seq_L
        self.good_mutations = [0] * self.seq_L
        self.mutation_score = [[]] * self.seq_L
        self.reverse_log = False
        self.current_mutations = []
        self.mutation_log = []
        self.last_update = time.time()
        self.last_graph = time.time()

        # Main loop:
        for self.step in range(self.N + 1):

            # random mutation at random position, also fix sequence constraint
            mut_start = time.time()
            if self.step != 0:
                seq_curr = self.mutate(seq)
                #if cfg.MATRIX_MODE == 'msa':
                #    for i in range(10): seq_curr = self.mutate(seq_curr)
            else: seq_curr = seq

            # Preprocess the sequence, relatively expensive operations (0.04s per pass)
            seq_curr = torch.from_numpy(seq_curr).long()
            model_input, msa1hot = prep_seq(seq_curr)
            muttime += time.time() - mut_start

            # probe effect of mutation
            nn_start = time.time()
            structure_predictions = self.structure_models( #run trRosettaEnsemble -> runs trrosetta n times
                model_input, use_n_models=cfg.n_models
            )
            nntime += time.time() - nn_start

            #loss functions
            loss_start = time.time()
            E_curr, metrics = self.loss(
                seq_curr, structure_predictions, msa1hot, track=True
            )
            losstime += time.time() - loss_start

            #scoring evaluation
            score_start = time.time()
            if E_curr < self.best_E:
                self.best_bkg = metrics["background_loss"]
                self.best_mtf = metrics["motif_loss"]
                self.best_site = metrics["site_loss"]
                self.best_distogram_distribution = structure_predictions['dist']

            seq, E = self.metropolis(seq, seq_curr, E_curr, E)

            E_tracker.append(v(E))
            delta_step_best = self.step - self.best_step
            scoretime += time.time() - score_start

            misc_start = time.time()
            if time.time() - self.last_update > cfg.report_interval or self.step == 0:
                fps = self.step / (time.time() - start_time)
                background_loss = metrics["background_loss"].cpu().detach().numpy()
                mtf_loss = metrics["motif_loss"].cpu().detach().numpy()
                site_loss = metrics["site_loss"].cpu().detach().numpy()

                b_bkg = self.best_bkg.cpu().detach().numpy()
                b_mtf = self.best_mtf.cpu().detach().numpy()
                b_site = self.best_site.cpu().detach().numpy()

                self.last_update = time.time()

                print(
                    f"Step {self.step:4d} / {self.N:4d} ({delta_step_best}) || "
                    f"beta: {self.beta:.1f}, "
                    f"mutations/s: {fps:.2f}, "
                    f"bad/good_accepts: {np.sum(self.bad_accepts[-100:])}/{np.sum(self.good_accepts[-100:])}",
                    flush=True,
                )
                print(f"STATS      || loss: {E_curr:.3f}, bkg: {background_loss:.3f}, mtf: {mtf_loss:.3f}, site: {site_loss:.3f}")
                print(f"BEST STATS || loss: {self.best_E:.3f}, bkg: {b_bkg:.3f}, mtf: {b_mtf:.3f}, site: {b_site:.3f}")

                #timings
                total_time = time.time() - start_time
                print(f"total time: {total_time:.1f}")
                print(f"mut time:   {muttime:.2f}\t| {100*muttime/total_time:.2f}%")
                print(f"nn time:    {nntime:.1f}\t| {100*nntime/total_time:.2f}%")
                print(f"loss time:  {losstime:.2f}\t| {100*losstime/total_time:.2f}%")
                print(f"score time: {scoretime:.1f}\t| {100*scoretime/total_time:.2f}%")
                print(f"misc time:  {misctime:.2f}\t| {100*misctime/total_time:.2f}%")

                plot_muts(self.gradient, self.results_dir / "gradient.jpg")
                plot_muts(self.good_mutations, self.results_dir / "mutations.jpg")

                if time.time() - self.last_graph > 1.8*cfg.report_interval or self.step == 0:
                    self.last_graph = time.time()

                    distogram = distogram_distribution_to_distogram(
                       self.best_distogram_distribution.cpu().detach().numpy()
                    )
                    plot_distogram(
                       distogram,
                       self.results_dir
                       / "distogram_evolution"
                       / f"{self.step:06d}_{self.best_E:.4f}.jpg",
                       clim=cfg.limits["dist"],
                    )

                    plot_progress(
                        E_tracker,
                        self.results_dir / "progress.jpg",
                        title=f"Optimization curve after {self.step} steps",
                    )
                    print(f"\n--- Current best: {self.best_sequence}")

            if (1 + self.step) % (cfg.report_interval//2) == 0:
                with open('control.txt', 'r') as reader:
                    lines = reader.readlines()
                    line = lines[0].strip()
                if line == "eval":
                    cmd = lines[1].strip()
                    print("EVALUATE")
                    print("cmd: " + cmd)
                    try:
                        eval(cmd.strip())
                    except:
                        print("ERROR")
                if line == "exit" or line == "break":
                    print("ending due to command")
                    break
                elif line == 'matrix':
                    if lines[1].strip() == 'n': cfg.MATRIX = False
                    else: cfg.MATRIX = True
                    print("MATRIX: " + str(cfg.MATRIX))
                elif line == 'gradient':
                    if lines[1].strip() == 'n': cfg.GRADIENT = False
                    else: cfg.GRADIENT = True
                elif line == "beta":
                    try:
                        value = float(lines[1].strip())
                        self.beta = value
                    except: print("input error")
                elif line == "bkw":
                    try:
                        value = float(lines[1].strip())
                        self.bkg_weight = value
                    except: print("input error")
                elif line == "aaw":
                    try:
                        value = float(lines[1].strip())
                        self.aa_weight = value
                    except: print("input error")
                elif line == "mw":
                    try:
                        value = float(lines[1].strip())
                        self.motif_weight = value
                    except: print("input error")
                elif line == "pause":
                    while True:
                        with open('control.txt', 'r') as reader:
                            line = reader.readlines()[0].strip()
                        if line != "pause":
                            break
                        else:
                            print("pause...")
                            time.sleep(30)

            if (delta_step_best > 300 and self.step > 1500 and self.step % 100 == 0) or (cfg.FAST and self.step > 499 and self.step % 50 == 0):
                std = np.std(np.array(E_tracker)[-1000:])
                n_std = std / np.array(E_tracker)[-1000:].mean()
                print("sd: " + str(std))
                if abs(std) < 0.005 or (cfg.FAST and abs(std) < 0.02):
                    break

            if self.step % self.M == 0 and self.step != 0:
                seq = self.fixup_MCMC(seq)

            misctime += time.time() - misc_start

        ########################################

        # Save final results before exiting:
        seq_curr = torch.from_numpy(
            aa2idx(self.best_sequence).reshape([1, self.seq_L])
        ).long()
        model_input, msa1hot = prep_seq(seq_curr)
        structure_predictions = self.structure_models(model_input)
        E_curr, self.best_metrics = self.loss(
            seq_curr, structure_predictions, msa1hot, track=True
        )

        #for key in self.best_metrics.keys():
        #    self.best_metrics[key] = v(metrics[key])

        TM_score_proxy = top_prob(structure_predictions['dist'], verbose=False)
        TM_score_proxy = TM_score_proxy[0]  # We're running with batch_size = 1

        self.best_metrics["TM_score_proxy"] = TM_score_proxy
        self.best_metrics["sequence"] = self.best_sequence
        self.best_metrics["motifweight"] = self.motif_weight
        self.best_metrics["motifmode"] = self.motifmode
        self.best_metrics["steps"] = self.step
        self.best_metrics["motifs"] = "\"" + str(self.motifs.copy()) + "\""

        # Dump distogram:
        best_distogram_distribution = structure_predictions['dist'].detach().cpu().numpy()
        distogram = distogram_distribution_to_distogram(best_distogram_distribution)
        plot_distogram(
            distogram,
            self.results_dir / f"result.jpg",
            clim=cfg.limits["dist"],
        )
        plot_progress(
            E_tracker,
            self.results_dir / "progress.jpg",
            title=f"Optimization curve after {self.step} steps",
        )

        # Write results to csv:
        with (self.results_dir / "result.csv").open("w") as f:
            for key, val in self.best_metrics.items():
                f.write(f"{key},{val}\n")

        with (self.results_dir / "mutation_log.csv").open("w") as f:
            for mut in self.mutation_log:
                f.write(f"{mut}\n")

        return self.best_metrics


import sys, os
import functools
import signal
import select
import socket
import numpy as np # CAUTION: Please use numpy version 1.15.4
import pickle
import time

from io import BytesIO
from pylab import figure,axes,pie,title,show
import time
from datetime import datetime
from multiprocessing import Process, Queue
import threading
import queue as queue
import interface
import struct
from PIL import Image

from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from PyQt5.QtQuick import *
from PyQt5.QtQml import *

#GLOBALS
headerStruct = struct.Struct('III')
MAX_BUFFER = 3 #since we are reading in image files, create a ring buffer and overwrite

class guiclient(QThread):
    sig_img = pyqtSignal(int,str, np.ndarray)
    sig_hist_val = pyqtSignal(int,int)
    sig_img_complete = pyqtSignal(str,int,int)# Time, width and height
    def __init__(self,server,port):
        QThread.__init__(self)
        super().__init__()
        print(server, port)
        self.sock = None
        self.win_mode = "" #to be used to determine what type of window, e.g. fourier, amplitude, phase, etc.
        self.exit = False
        self.maxlen = 150995023
        self.counter = 0
        self.histogram =0
        self.bins = 0
        self.clientMontiorThread = None
        self.server = server
        self.port = port
        self.reconst = False
        self.wavelength = 0
        self.prop_distance = 0
        self.performance_val = 1
        self.quit = False
        self.outdata = None
        self.init = False
        self.outdata_RGB = None
        self.enable_histogram = False

        self.z = np.zeros((2448,2050),dtype=np.uint8)
        self.w = None
        self.h = None
        self.l = None
        self.msgid = None
        self.srcid = None
        self.totalbytes = None
        self.msg = None
        self.meta = None
        self.offset = None

        self.ndim_struct = None
        self.ndimsize = None
        self.ndim = None

        self.dimStruct = None
        self.dimsize = None
        self.dimensions = None

        self.dtype = None

        print("*** DISPLAY WINDOW CREATED ***")

        for sig in [signal.SIGTERM, signal.SIGINT, signal.SIGHUP, signal.SIGQUIT]:
            signal.signal(sig, self.signal_handler)

    def set_mode(self,mode):
        if(mode == "reconst"):
           self.reconst = True
        else:
           self.reconst = False

    def signal_handler(self, signal, frame):
        print("DHMx: Gui Thread exiting...")
        self.exit = True
        sys.exit()

    def restore_frame(self, data, meta, z):
        w, h, compval, val, size, actualsize,ts, gain, ccdtemp = meta

        dtidx = 0
        for i in range(0, w*h,compval):
            z[i] = data[dtidx];
            dtidx += 1

    # x and y values are flipped in numpy array (in dhmx.py it is transmitted as x,y).
    def get_pixel_val(self,y,x):
       # Check if the image sending has been initialized by its first frame to sample
       if(self.init != False ):
          # Check to make sure the mouse is within bounds.  Some images may not be exactly 2048x2048
          if((np.size(self.outdata,0) >= x and np.size(self.outdata,1) >= y)):
             if(self.reconst and self.win_mode != 'fourier'):
                return int(self.outdata[x-1,y-1,self.prop_distance,self.wavelength])
             else:
                return int(self.outdata[x-1,y-1])
       else:
          return 0

    # This will launch and run the Display Thread as a Qt thread
    def unpack_message(self,msg):
            self.msg = msg
            try:
               self.msgid, self.srcid, self.totalbytes = headerStruct.unpack(self.msg[0:struct.calcsize(headerStruct.format)])
            except:
               print("DHMx Diplsay: Did not receive proper header.  Header missing or incomplete.")
               return -10
            self.meta = (self.msgid, self.srcid, self.totalbytes)
            self.offset = struct.calcsize(headerStruct.format)

            self.ndim_struct = struct.Struct('H')
            self.ndimsize = struct.calcsize(self.ndim_struct.format)
            self.ndim = self.ndim_struct.unpack(self.msg[self.offset:self.offset + self.ndimsize])[0]

            self.dimStruct = struct.Struct('H'*int(self.ndim))
            self.dimsize = struct.calcsize(self.dimStruct.format)
            self.dimensions = self.dimStruct.unpack(self.msg[self.offset + self.ndimsize:self.offset + self.ndimsize + self.dimsize])

            self.offset = self.offset + self.ndimsize + self.dimsize

            if self.srcid == interface.SRCID_IMAGE_FOURIER:
                self.dtype = np.uint8
                self.h, self.w = self.dimensions
            elif self.srcid == interface.SRCID_IMAGE_RAW:
                self.dtype = np.uint8
                self.h, self.w = self.dimensions
            else:
                self.dtype = np.uint8
                self.h, self.w, self.z, self.l = self.dimensions

            if(self.srcid == interface.SRCID_IMAGE_AMPLITUDE and self.win_mode == "amplitude") or\
              (self.srcid == interface.SRCID_IMAGE_PHASE and self.win_mode == "phase") or\
              (self.srcid == interface.SRCID_IMAGE_FOURIER and self.win_mode == "fourier") or\
              (self.srcid == interface.SRCID_IMAGE_RAW and self.win_mode == "raw") or\
              (self.srcid == interface.SRCID_IMAGE_INTENSITY and self.win_mode == "intensity"):

                 # Grab data from the DHM server...
                 self.outdata = np.frombuffer(self.msg[self.offset:self.offset+(functools.reduce(lambda x,y: x*y, self.dimensions)*np.dtype(self.dtype).itemsize)], dtype=self.dtype).reshape(self.dimensions)

                 # make the array writable
                 try:
                     self.outdata.setflags(write=1)
                 except:
                     pass

                 #rescaled for performance increase of upwards of 40+%
                 self.outdata = self.outdata[::self.performance_val, ::self.performance_val]

                 # Rescale image
                 if(self.outdata.max() != 0):
                    self.outdata *= int(255.0/self.outdata.max())

                 # If the user selects the histogram button and enables the histogram, compute the bins necessary to display
                 # NOTE: Histogram computations in real time are CPU costly, so it is best to only ue this if/when needed.
                 if(self.enable_histogram):
                    # calculate the histogram on the grayscale 'outdata' object from 0-255
                    # Emit the position of the histogram back to Qt/QML to update the histogram
                    if(self.reconst and self.win_mode != 'fourier'):
                       self.histogram,self.bins = np.histogram(self.outdata[:,:,self.prop_distance,self.wavelength],bins=np.arange(0,256,1))
                    else:
                       self.histogram,self.bins = np.histogram(self.outdata,bins=np.arange(0,256,1))
                    for i in range(255):
                       self.sig_hist_val.emit(i,self.histogram[i])

                 # Create an RGB version of the received image to display absolute minimums and maximums
                 if(self.reconst and self.win_mode != 'fourier'):
                    self.outdata_RGB = np.stack((self.outdata[:,:,self.prop_distance,self.wavelength],)*3, axis=-1)
                    self.outdata_RGB[self.outdata[:,:,self.prop_distance,self.wavelength] == 255] = [255,0,0]
                    self.outdata_RGB[self.outdata[:,:,self.prop_distance,self.wavelength] == 0] = [0,0,255]
                 else:
                    self.outdata_RGB = np.stack((self.outdata,)*3, axis=-1)
                    self.outdata_RGB[self.outdata == 255] = [255,0,0]
                    self.outdata_RGB[self.outdata == 0] = [0,0,255]

                 # Create an image timestamp in UTC format
                 current_time = datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]

                 # Emit to Qt/QML that the processed image is ready to be displayed
                 # The signal will emit a UTC time, the camera width and camera height
                 self.sig_img_complete.emit(current_time,self.w,self.h)
                 self.init = True

    def set_win(self,mode):
        self.win_mode = mode

    # gets index 0-255 and returns a value associated with the luma range
    # of 0 (absolute black) to 255 (absolute white)
    def get_histogram_val(self, pos):
        return self.bins[pos]

    def shutdown(self):
        self.clientMontiorThread.join()
        self.clientMontiorThread.close()
        print("Done")
        self.terminate()

    def recvall(self,sock,n):
       data = b''
       while len(data) < n:
          packet = sock.recv(n - len(data))
          if not packet:
             return None
          data += packet
       return data

    def run(self):
        ### Continous receive of data
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((self.server, self.port))
        self.readfds = [self.sock]

        ### Start Display Thread
        #self.displaythread.start()
        length = None
        buf = b''
        data = b''
        msg=b''
        meta = None
        totalbytes = 0

        while not self.quit:
            infds, outfds, errfds = select.select(self.readfds, [], [])

            if not (infds or outfds or errfds):
                continue
            if self.exit: break

            for s in infds:
                if s is self.sock:
                    ### Get as much data as we can
                    #packet = self.sock.recv(150995023)

                    # The data size is 2048x2048 * 3
                    packet = self.sock.recv(12582912)
                    if not packet:
                        self.exit = True
                        break
                    data += packet
                    datalen = len(data)
                    if meta is None and datalen > struct.calcsize(headerStruct.format):
                        try:
                            msg_id, srcid, totalbytes = headerStruct.unpack(data[0:struct.calcsize(headerStruct.format)])
                        except:
                            print("DHMx Diplsay: Did not receive proper header.  Header missing or incomplete.")
                            return -10
                        totalbytes += struct.calcsize(headerStruct.format)
                        self.meta = (msg_id, srcid)
                        if datalen >= totalbytes:
                            msg = data[:totalbytes]
                            data = data[totalbytes:]
                            meta = None
                            totalbytes = 0
                            self.unpack_message(msg)
                    else:
                        if datalen < totalbytes:
                            continue
                        msg = data[:totalbytes]
                        data = data[totalbytes:]
                        meta = None
                        totalbytes = 0
                        self.unpack_message(msg)

            if self.exit: break
        self.sock.close()


from flask import *

import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array

app = Flask(__name__)

MODEL_PATH ='Gender-Model.h5'
model = load_model(MODEL_PATH)

@app.route('/', methods=['GET'])
def home():
    return render_template('index.html')

@app.route('/', methods=['POST'])
def predict():
    imagefile = request.files['imagefile']
    image_path = "./images/" + imagefile.filename
    imagefile.save(image_path)

    image = load_img(image_path, target_size=(250,250))
    image = img_to_array(image)
    image = np.expand_dims(image, axis=0)

    preds = model.predict(image)

    if preds[0][0] == 1:
         classification = "He is a Male 🏃"
    else:
         classification = "She is a Female 💃"

    return render_template('index.html', prediction = classification)

if __name__ == '__main__':
    app.run(debug=True, port=3000)

Gender Classification using ConvNets/app.py


import face_recognition
import os
import cv2
from os import listdir, remove, mkdir
import pickle
import numpy as np
from operator import add
import json
from shutil import copyfile, rmtree
from utils import improve_contrast_image_using_clahe, image_resize, image_size_limit, generate_weights
from app import start_app

SRC_IMAGE_PATH = 'images/'
DST_IMAGE_PATH = 'static/images/'
ENCODINGS = 'encodings.dat' # just a caching file
TOLERANCE = 0.6 # 0.6 is default
MODEL = 'cnn' # use 'hog' if you dont have CUDA

def load_face_encodings(image_path, encodings_path):
    """Load face encodings from cache or from files."""

    face_id = 0

    # checking for cache
    if encodings_path in os.listdir():
        print('load encodings from cache')
        with open(encodings_path, 'rb') as f:
            tmp = pickle.load(f)
            return tmp[0], tmp[1]
    else:
        print('load encodings from files')
        encodings = []
        crop_sizes = []
        files = os.listdir(image_path)
        max_len = len(max(files, key=len))

        for file in files:

            print('loading', file, end=" "*max_len + '\r', flush=True)

            # load image-file
            image = face_recognition.load_image_file(f'{image_path}/{file}')

            # preprocess image
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            image = improve_contrast_image_using_clahe(image)
            image = image_size_limit(image, 2000) # change if out of memory

            # find faces in image
            locations = face_recognition.face_locations(image, model=MODEL, number_of_times_to_upsample=1)
            encoding = face_recognition.face_encodings(image, locations, num_jitters=10, model='large')
            encodings += encoding

            for face_location in locations:

                # crop face out of the image
                top_left = (face_location[3], face_location[0])
                bottom_right = (face_location[1], face_location[2])
                crop = image[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]

                # get size of the face
                h, w, _ = crop.shape
                crop_size = h*w
                crop_sizes.append(crop_size)

                # save face in /crops
                crop = image_resize(crop, height=256)
                cv2.imwrite(f'crops/{face_id}.png', crop)
                face_id += 1

        temp = []
        for i in range(len(encodings)):
            temp.append([crop_sizes[i], encodings[i], i])
        temp = sorted(temp, key=lambda l:l[0], reverse=True)
        encs = []
        indexes = []
        for i in range(len(temp)):
            encs.append(temp[i][1])
            indexes.append(temp[i][2])

        print('all images loaded')
        with open(encodings_path, 'wb') as f:
            pickle.dump([encs, indexes], f)
        return encs, indexes

def average_encoding(encodings):
    """Return average encoding from a list of encoding."""
    ret = [0]*128
    weights = generate_weights(len(encodings))
    for i in range(len(encodings)):
        for j in range(len(encodings[i])):
            ret[j] += encodings[i][j]*weights[i]
    return ret

def average_encodings(encodings):
    """Return a list of average encodings from a 2d-list of encodings."""
    ret = []
    for encoding in encodings:
        if len(encoding) > 1:
            tmp = average_encoding(encoding)
            ret.append(tmp)
        else:
            ret.append(encoding[0])
    return np.array(ret).tolist()

def face_grouping(encodings, indexes):
    """Face grouping algo."""
    known = [[encodings[0]]]
    known_orgs = [[indexes[0]]]
    for i in range(1, len(encodings)):
        tmp = face_recognition.face_distance(average_encodings(known), encodings[i])
        min_dist = tmp.min()
        if min_dist <= TOLERANCE:
            index = np.where(tmp==min_dist)[0][0]
            known[index].append(encodings[i])
            known_orgs[index].append(indexes[i])
        else:
            known.append([])
            known[-1].append(encodings[i])
            known_orgs.append([])
            known_orgs[-1].append(indexes[i])
    return known, known_orgs

def generate_json(known, known_orgs, DST_IMAGE_PATH):
    """Generate data.json."""
    nodes = [{'name': str(i)} for i in range(len(known))]
    for i in range(len(known_orgs)):
        copyfile('crops/' + str(known_orgs[i][0]) + '.png', DST_IMAGE_PATH + str(i) + '.png')
    links = []
    for i in range(len(known)-1):
        for j in range(i+1, len(known)):
            links.append({'source': i,
                          'target': j,
                          'distance': face_recognition.face_distance([average_encoding(known[i])], known[j][0])[0]})
    data = {'nodes': nodes, 'links': links}
    with open('data.json', 'w') as f:
        json.dump(data, f)

def drop_cache():
    """Drop all cached files."""
    rmtree('crops')
    mkdir('crops')
    rmtree('static/images')
    mkdir('static/images')
    if 'encodings.dat' in listdir():
        remove('encodings.dat')
    if 'data.json' in listdir():
        remove('data.json')

if __name__ == '__main__':

    # drop cache
    if input('Do you want to delete the cache? (YES/no)').lower() in ['yes', 'y', '']:
        drop_cache()

    # load encodings
    encodings, indexes = load_face_encodings(SRC_IMAGE_PATH, ENCODINGS)
    print('found', len(encodings), 'faces')

    # group faces
    known, known_orgs = face_grouping(encodings, indexes)

    # generating JSON
    generate_json(known, known_orgs, DST_IMAGE_PATH)

    start_app()


import pickle
from random import random
import torch
import time
import pyautogui
import gc
import numpy as np
from threading import Thread

from tqdm import tqdm

import environment
from trainer import QTrainer, RainbowTrainer, TaikoTrainer

torch.set_printoptions(sci_mode=False)

pyautogui.MINIMUM_DURATION = 0.0
pyautogui.MINIMUM_SLEEP = 0.0
pyautogui.PAUSE = 0.0

BATCH_SIZE = 32
LEARNING_RATE = 0.00005  # Double DQN: 0.00025, Prioritized Replay -> Double DQN / 4
GAMMA = 0.999
MAX_STEPS = 25000
WIDTH = 878
HEIGHT = 600
STACK_SIZE = 4

DISCRETE_FACTOR = 10
X_DISCRETE = 685 // DISCRETE_FACTOR + 1
Y_DISCRETE = (560 - 54) // DISCRETE_FACTOR + 1

PIXEL_SKIP = 4

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def print_infos():
    print('Discretized x: ' + str(X_DISCRETE))
    print('Discretized y: ' + str(Y_DISCRETE))
    print('Action dim: ' + str(X_DISCRETE * Y_DISCRETE * 4))
    print('X_MAX = ' + str(145 + (X_DISCRETE - 1) * DISCRETE_FACTOR))
    print('YMAX = ' + str(54 + 26 + (Y_DISCRETE - 1) * DISCRETE_FACTOR))

def log_episodes(i, q_trainer, steps, delta_t, k):  # TODO : Make a logging utils and use log function
    if i <= 1 or i % 30 == 0:
        print(str(steps) + ' steps in ' + str(delta_t) + ' s.' + '  -->   %.4f' % (steps / delta_t) + ' steps / s')

    if i > 0 and i % 50 == 0:
        q_trainer.avg_reward_plotter.fit()
    q_trainer.avg_reward_plotter.show()

    if i > 0 and i % 250 == 0:
        q_trainer.plotter.fig.savefig('average_loss' + str(i) + '.png')
        q_trainer.avg_reward_plotter.fig.savefig('episode_reward' + str(i) + '.png')

    if i % 30 == 0:
        print("Episode %d" % (i+1))
        print("Time step %d" % k)
        print("Schedule: %f" % q_trainer.scheduler.value(k))

def trainQNetwork(episode_nb, learning_rate, batch_size=BATCH_SIZE, load_weights=None, save_name='tests',
                  beatmap_name=None, star=1, evaluation=False, human_off_policy=False, load_memory=None, no_fail=False,
                  initial_p=1.0, end_p=0.05, decay_p=2000000, target_update=30000, init_k=0, min_experience=25000):
    print_infos()
    if evaluation:
        learning_rate = 0.0

    env = environment.OsuEnv(X_DISCRETE, Y_DISCRETE, DISCRETE_FACTOR, WIDTH, HEIGHT, STACK_SIZE, star=star,
                             beatmap_name=beatmap_name, no_fail=no_fail, skip_pixels=PIXEL_SKIP)
    q_trainer = QTrainer(env, batch_size=batch_size, lr=learning_rate, gamma=GAMMA, initial_p=initial_p, end_p=end_p,
                         decay_p=decay_p, load_weights=load_weights, load_memory=load_memory,
                         min_experience=min_experience, gradient_clipping_norm=10.0)

    episode_average_reward = 0.0
    k = init_k
    reward = torch.tensor(0.0, device=device)
    for i in range(episode_nb):
        controls_state, state = env.reset(reward)
        env.launch_episode(reward)

        episode_reward = 0.0
        thread = None
        start = time.time()
        for steps in range(MAX_STEPS):
            previous_t = time.time()
            k += 1
            with torch.no_grad():
                if not human_off_policy:
                    if evaluation:  # Choose greedy policy if tests
                        action = q_trainer.select_action(state, controls_state)
                    else:  # Else choose an epsilon greedy policy with decaying epsilon
                        if k > min_experience:
                            sample = random()
                            if sample > q_trainer.scheduler.value(k):
                                action = q_trainer.select_action(state, controls_state)
                            else:
                                action = q_trainer.random_action(X_DISCRETE, Y_DISCRETE)
                        else:
                            action = q_trainer.random_action(X_DISCRETE, Y_DISCRETE)

                    new_state, new_controls_state, reward, done = env.step(action, steps)
                else:
                    action, new_state, new_controls_state, reward, done = env.observe(steps)

                if done:
                    new_state = None
                else:
                    th = Thread(target=q_trainer.memory.push,
                                args=(state, action, reward, new_state, controls_state, new_controls_state))
                    th.start()

            if thread is not None:
                thread.join()
            thread = Thread(target=q_trainer.optimize)
            thread.start()

            state = new_state
            controls_state = new_controls_state
            episode_reward += reward

            #busy_wait(freq=12.0, previous_t=previous_t)
            if done:
                break
        end = time.time()
        delta_t = end - start

        gc.collect()

        episode_average_reward += episode_reward
        q_trainer.avg_reward_plotter.step(episode_reward)
        log_episodes(i, q_trainer, steps, delta_t, k)

        if k % target_update == 0:
            q_trainer.target_q_network.load_state_dict(q_trainer.q_network.state_dict())

        if i % 30 == 0 and i > 0:
            print('Mean reward over last 30 episodes: ')
            print(episode_average_reward / 30)

            episode_average_reward = 0.0
            if beatmap_name is not None:
                tmp = beatmap_name + save_name
            else:
                tmp = save_name
            q_trainer.save_model(tmp, num=i)

    if (episode_nb - 1) % 15 != 0:
        if beatmap_name is not None:
            tmp = beatmap_name + save_name
        else:
            tmp = save_name
        q_trainer.save_model(tmp, num=episode_nb - 1)
        q_trainer.plotter.fig.savefig('average_loss' + str(episode_nb-1) + '.png')
        q_trainer.avg_reward_plotter.fig.savefig('average_reward' + str(episode_nb-1) + '.png')
    print('k: %d' % k)
    env.stop()

def RainbowManiaTrain(lr=0.0000625, batch_size=32, gamma=0.999, omega=0.5, beta=0.4, sigma=0.1, eps=1.5e-4, n=3, atoms=51,
                      max_timestep=int(5e7), learn_start=80000, stack_size=4, norm_clip=10.0, save_freq=50000,
                      model_save_path='weights/Rainbow_test', memory_save_path='weights/memory.zip', target_update_freq=80000,
                      star=4, beatmap_name=None, width=380, height=600, skip_pixels=4, num_actions=128, no_fail=False,
                      load_weights=None, load_memory=None, Vmin=-10, Vmax=10, resume_start=0, load_stats=None,
                      save_stats='./stats.pkl', learning_freq=1, load_optimizer=None, optimizer_path='weights/opti.pt',
                      evaluation=False, data_efficient=False):

    priority_weight_increase = (1 - beta) / (max_timestep - learn_start - resume_start)

    env = environment.ManiaEnv(height=height, width=width, stack_size=stack_size, star=star, beatmap_name=beatmap_name,
                               num_actions=num_actions, skip_pixels=skip_pixels, no_fail=no_fail)
    trainer = RainbowTrainer(env, batch_size=batch_size, lr=lr, gamma=gamma, omega=omega, beta=beta, sigma=sigma, n=n,
                             eps=eps, atoms=atoms, norm_clip=norm_clip, load_weights=load_weights, load_memory=load_memory,
                             Vmin=Vmin, Vmax=Vmax, load_optimizer=load_optimizer, data_efficient=data_efficient)
    if load_stats is None:
        stat = {'episode_reward': []}
    else:
        with open(load_stats, 'rb') as f:
            stat = pickle.load(f)

    reward = 0.0
    need_save = False
    done = True
    count = 0
    thread = None

    for t in tqdm(range(resume_start, max_timestep), desc="Timestep", unit='step', unit_scale=True):
        if done:
            if t > resume_start:
                count += 1
                trainer.avg_reward_plotter.step(episode_reward)
                trainer.avg_reward_plotter.show()
                trainer.plotter.show()
                stat['episode_reward'].append(episode_reward.item())
                if count % 100 == 0:
                    print('   -  Mean reward over last 100 episodes: %.4f' % np.array(stat['episode_reward'][max(-100, -len(stat['episode_reward'])):]).mean())
            if need_save:
                trainer.save(model_save_path + str(t) + ".pt", memory_save_path, optimizer_path)
                with open(save_stats, 'wb') as f:
                    pickle.dump(stat, f, protocol=4)
                need_save = False

            episode_reward = 0.0
            gc.collect()
            state = env.reset(reward)
            env.launch_episode(reward)
        trainer.reset_noise()
        action = trainer.select_action(state)
        next_state, reward, done = env.step(action)
        reward = max(min(reward, 1.0), -1.0)  # Reward clipping
        episode_reward += reward
        if env.episode_counter > 0:  # Skip first episode because of latency issues
            trainer.memory.append(state[-1], action, reward, done)

        if t >= learn_start and t % learning_freq == 0 and not evaluation:
            trainer.memory.priority_weight = min(trainer.memory.priority_weight + priority_weight_increase, 1)
            if thread is not None:
                thread.join()
            thread = Thread(target=trainer.optimize)
            thread.start()

        if t % target_update_freq == 0 and t > 0:
            trainer.update_target_net()

        if t % save_freq == 0 and t > 0:
            need_save = True

def TaikoTrain(lr=0.00008, batch_size=32, stack_size=1, skip_pixels=4, save_freq=20000, episode_nb=5,
                target_update_freq=10000, star=None, beatmap_name=None, min_experience=10000, root_dir='./weights',
                evaluation=False):

    env = environment.TaikoEnv(stack_size=stack_size, star=star, beatmap_name=beatmap_name, skip_pixels=skip_pixels)
    tt = TaikoTrainer(env, batch_size=batch_size, lr=lr, gamma=GAMMA, root_dir=root_dir, min_experience=min_experience, norm_clip=10.0)

    need_update = False
    need_save = False
    for episode in range(episode_nb):
        state = env.reset()
        env.launch_episode()
        episode_reward = 0.0
        for steps in range(MAX_STEPS):
            if evaluation:
                action = tt.select_action(state)
            else:
                action = tt.select_explo_action(state)

            new_state, reward, done = env.step(action)
            episode_reward += reward

            if not done:
                th = Thread(target=tt.memory.push,
                            args=(state, action, reward, new_state))
                th.start()
            else:
                break

            state = new_state
            # These boolean allow the program to wait for the end of the episode before performing the updates or the save to avoid latency while the agent is playing
            if tt.steps_done % target_update_freq == 0:
                need_update = True
            if tt.steps_done % save_freq == 0:
                need_save = True

        for _ in range(steps):
            tt.optimize()

        if need_update:
            tt.update_target()
            need_update = False
        if need_save:
            tt.save()
            need_save = False
        print(steps)
        tt.avg_reward_plotter.step(episode_reward)
        tt.avg_reward_plotter.show()

    tt.stop()

if __name__ == '__main__':
    '''
    weights_path = './weights/q_net__21-12-2020-14.pt'
    memory_path = './memory.pck'
    save_name = '_25-12-2020-'
    trainQNetwork(50, LEARNING_RATE, evaluation=False, load_weights=None, beatmap_name="sink", star=7,
                  save_name=save_name, batch_size=BATCH_SIZE, human_off_policy=False, no_fail=True,
                  initial_p=1.0, end_p=0.05, decay_p=4000000, target_update=30000, init_k=0, min_experience=50)
    '''
    '''
    RainbowManiaTrain(star=3, beatmap_name="bongo", num_actions=2**4, model_save_path="weights/Taiko_Bongo-31-05-2021_3stars",
                      learn_start=1600, load_weights=None, load_memory=None, batch_size=32, max_timestep=int(1e6),
                      memory_save_path='./weights/memory28-03-2021.zip', Vmin=-1, Vmax=10, resume_start=0, target_update_freq=5000,
                      load_stats=None, save_freq=5000, save_stats='./stats/stats-28-03-2021.pkl', learning_freq=1, lr=0.0001,
                      load_optimizer=None, optimizer_path='./weights/opti.pt', evaluation=False, n=20, data_efficient=True)
    '''

    TaikoTrain(root_dir='./weights/Taiko/2021-09-05_15/29/', evaluation=True, episode_nb=1500, min_experience=1500, save_freq=5000, target_update_freq=2000, batch_size=16, lr=0.0002)


import colorsys
from timeit import default_timer as timer

import numpy as np
from keras import backend as K
from keras.models import load_model
from keras.layers import Input
from PIL import Image, ImageFont, ImageDraw
from yolo3.model import yolo_eval, yolo_body, tiny_yolo_body,box_iou
from yolo3.utils import letterbox_image
import os
from keras.utils import multi_gpu_model

class YOLO(object):
    _defaults = {
        "model_path": 'model_data/yolo.h5',
        "anchors_path": 'model_data/yolo_anchors.txt',
        "classes_path": 'model_data/coco_classes.txt',
        "score": 0.3,
        "iou": 0.45,
        "model_image_size": (416, 416),
        "gpu_num": 1,
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)  # set up default values
        self.__dict__.update(kwargs)  # and update with user overrides
        self.class_names = self._get_class()
        self.anchors = self._get_anchors()
        self.sess = K.get_session()
        self.generate()
        # Initialize the parameters
        self.confThreshold = 0.5  # Confidence threshold
        self.nmsThreshold = 0.4  # Non-maximum suppression threshold
        self.inpWidth = 416  # Width of network's input image
        self.inpHeight = 416  # Height of network's input image

    def _get_class(self):
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    def _get_anchors(self):
        anchors_path = os.path.expanduser(self.anchors_path)
        with open(anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape(-1, 2)

    def generate(self):
        model_path = os.path.expanduser(self.model_path)
        assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'

        # Load model, or construct model and load weights.
        num_anchors = len(self.anchors)
        num_classes = len(self.class_names)
        is_tiny_version = num_anchors == 6  # default setting
        try:
            self.yolo_model = load_model(model_path, compile=False)
        except:
            self.yolo_model = tiny_yolo_body(Input(shape=(None, None, 3)), num_anchors // 2, num_classes) \
                if is_tiny_version else yolo_body(Input(shape=(None, None, 3)), num_anchors // 3, num_classes)
            self.yolo_model.load_weights(self.model_path)  # make sure model, anchors and classes match
        else:
            assert self.yolo_model.layers[-1].output_shape[-1] == \
                   num_anchors / len(self.yolo_model.output) * (num_classes + 5), \
                'Mismatch between model and given anchor and class sizes'

        print('{} model, anchors, and classes loaded.'.format(model_path))

        # Generate colors for drawing bounding boxes.
        hsv_tuples = [(x / len(self.class_names), 1., 1.)
                      for x in range(len(self.class_names))]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                self.colors))
        np.random.seed(10101)  # Fixed seed for consistent colors across runs.
        np.random.shuffle(self.colors)  # Shuffle colors to decorrelate adjacent classes.
        np.random.seed(None)  # Reset seed to default.

        if self.gpu_num >= 2:
            self.yolo_model = multi_gpu_model(self.yolo_model, gpus=self.gpu_num)

    def detect_image(self, image):
        start = timer()

        if self.model_image_size != (None, None):
            assert self.model_image_size[0] % 32 == 0, 'Multiples of 32 required'
            assert self.model_image_size[1] % 32 == 0, 'Multiples of 32 required'
            boxed_image = letterbox_image(image, tuple(reversed(self.model_image_size)))
        else:
            new_image_size = (image.width - (image.width % 32),
                              image.height - (image.height % 32))
            boxed_image = letterbox_image(image, new_image_size)
        image_data = np.array(boxed_image, dtype='float32')

        print(image_data.shape)
        image_data /= 255.
        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.

        prediction = self.yolo_model.predict(image_data)
        boxes = self.postProcess(prediction,self.anchors,len(self.class_names),(image.size[1],image.size[0]))
        print('Found {} boxes for {}'.format(len(boxes), 'img'))

        font = ImageFont.truetype(font='font/FiraMono-Medium.otf',
                                  size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
        thickness = (image.size[0] + image.size[1]) // 300

        for top,left,bottom,right,box_classes,box_class_scores in boxes:
            box_classes=int(box_classes)
            predicted_class = self.class_names[box_classes]
            top = max(0, np.floor(top + 0.5).astype('int32'))
            left = max(0, np.floor(left + 0.5).astype('int32'))
            bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
            right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
            label = '{} {:.2f}'.format(predicted_class, box_class_scores)
            draw = ImageDraw.Draw(image)
            label_size = draw.textsize(label, font)
            print(label, (left, top), (right, bottom))

            if top - label_size[1] >= 0:
                text_origin = np.array([left, top - label_size[1]])
            else:
                text_origin = np.array([left, top + 1])

            # My kingdom for a good redistributable image drawing library.
            for i in range(thickness):
                draw.rectangle(
                    [left + i, top + i, right - i, bottom - i],
                    outline=self.colors[box_classes])
            draw.rectangle(
                [tuple(text_origin), tuple(text_origin + label_size)],
                fill=self.colors[box_classes])
            draw.text(text_origin, label, fill=(0, 0, 0), font=font)
            del draw

        end = timer()
        print(end - start)
        return image

    def sigmoid(self,x):
        return 1 / (1 + np.exp(-x))
    # Remove the bounding boxes with low confidence using non-maxima suppression

    def postProcess(self,feats,anchors, num_classes, image_shape):
        """Convert final layer features to bounding box parameters."""
        # Reshape to batch, height, width, num_anchors, box_params.
        anchor_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
        result = []
        boxes_list = []
        box_scores_list = []
        for l in range(len(feats)):
            anchor = anchors[anchor_mask[l]]
            anchors_tensor = np.reshape(anchor, [1, 1, 1, len(anchor), 2])
            feat = feats[l]
            grid_shape = np.shape(feat)[1:3] # height, width
            grid_y = np.tile(np.reshape(np.arange(0, stop=grid_shape[0]), [-1, 1, 1, 1]),
                            [1, grid_shape[1], 1, 1])
            grid_x = np.tile(np.reshape(np.arange(0, stop=grid_shape[1]), [1, -1, 1, 1]),
                            [grid_shape[0], 1, 1, 1])
            grid = np.concatenate((grid_x, grid_y),axis=-1)
            # grid = np.concatenate((grid_x, grid_y))
            feat = np.reshape(
                feat, [-1, grid_shape[0], grid_shape[1], 3, 80 + 5])

            box_confidence = self.sigmoid(feat[..., 4:5])
            box_class_probs = self.sigmoid(feat[..., 5:])

            input_shape = np.array([416,416])

            # Adjust preditions to each spatial grid point and anchor size.
            box_xy = (self.sigmoid(feat[..., :2]) + grid) / grid_shape[::-1]
            box_wh = np.exp(feat[..., 2:4]) * anchors_tensor / input_shape[::-1]

            box_yx = box_xy[..., ::-1]
            box_hw = box_wh[..., ::-1]
            new_shape = np.round(np.multiply(image_shape,min(input_shape / image_shape)))
            offset = (input_shape - new_shape) / 2. / input_shape
            scale = input_shape / new_shape
            box_yx = np.multiply((box_yx - offset), scale)
            box_hw *= scale

            box_mins = box_yx - (box_hw / 2.)
            box_maxes = box_yx + (box_hw / 2.)
            boxes = np.concatenate((
                box_mins[..., 0:1],  # y_min
                box_mins[..., 1:2],  # x_min
                box_maxes[..., 0:1],  # y_max
                box_maxes[..., 1:2]  # x_max
            ),axis=-1)

            # Scale boxes back to original image shape.
            boxes = np.multiply(boxes, np.concatenate((image_shape, image_shape)))
            box_scores = np.multiply(box_confidence,box_class_probs)

            #NMS
            boxes = np.reshape(boxes, [-1, 4])
            box_scores = np.reshape(box_scores, [-1, num_classes])
            boxes_list.append(boxes)
            box_scores_list.append(box_scores)
        boxes_list = np.concatenate(boxes_list, axis=0)
        box_scores_list = np.concatenate(box_scores_list, axis=0)
        for c in range(num_classes):
            scores = box_scores_list[:,c]
            boxes_class = np.insert(boxes_list,4, c,1)
            boxes_class = np.insert(boxes_class,5, scores,1)
            mask = scores>self.confThreshold
            boxes_class = boxes_class[mask]
            if len(boxes_class)>0:
                boxes_class=sorted(boxes_class,key=lambda x:x[-1],reverse=True)
                iou_list = np.array(boxes_class)

                while len(iou_list)>0:
                    candidate_box= iou_list[0]
                    result.append(candidate_box)
                    if len(iou_list)==1:
                        break
                    iou_list = iou_list[1:]

                    b1_mins = iou_list[:,:2]
                    b2_mins = candidate_box[:2]
                    b1_maxs = iou_list[:,2:4]
                    b2_maxs = candidate_box[2:4]
                    intersect_mins = np.maximum(b1_mins, b2_mins)
                    intersect_maxes = np.minimum(b1_maxs, b2_maxs)
                    intersect_wh = np.maximum(intersect_maxes - intersect_mins, 0.)
                    intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
                    b1_area = (b1_maxs[:,0] - b1_mins[:,0]) * (b1_maxs[:,1] - b1_mins[:,1])
                    b2_area = (b2_maxs[0] - b2_mins[0]) * (b2_maxs[1] - b2_mins[1])
                    iou = intersect_area / (b1_area + b2_area - intersect_area)
                    iou_mask = iou<=self.iou
                    # iou_mask =np.concatenate(iou_mask,axis=0)
                    iou_list = iou_list[iou_mask]
            #Without nms
            # boxes_class = np.concatenate((boxes,box_scores),axis=-1)[0]
            # for i in range(len(boxes_class)):
            #     grid_y_list = boxes_class[i]
            #     for j in range(len(grid_y_list)):
            #         grid = grid_y_list[j]
            #         for k in range(len(grid)):
            #             box = grid[k]
            #             for m in range(4,len(box)):
            #                 probability = box[m]
            #                 if probability>self.confThreshold:
            #                     result.append((box[0],box[1],box[2],box[3],m-4,probability))
        return result

if __name__ == '__main__':
    img = '.\TestData\\soccer_foul.jpg'
    image = Image.open(img)
    yolo = YOLO()
    r_image = yolo.detect_image(image)
    r_image.show()


from re import Match
from natsort import natsorted
import numpy as np
import transforms3d
import os
import re

from common.Util import Util
from common.ymap.Flag import Flag
from common.ymap.Ymap import Ymap
from common.ytyp.YtypItem import YtypItem
from common.ytyp.YtypParser import YtypParser

class Sanitizer:
    identityQuaternion = [1, -0, -0, -0]

    inputDir: str
    outputDir: str
    ytypItems: dict[str, YtypItem]
    lowercaseYtypItems: dict[str, str]
    fixedArchetypeNames: set[str]

    def __init__(self, inputDir: str, outputDir: str):
        self.inputDir = inputDir
        self.outputDir = outputDir

    def run(self):
        print("running sanitizer...")
        self.createOutputDir()
        self.readYtypItems()
        self.processFiles()
        self.copyOthers()
        print("sanitizer DONE")

    def createOutputDir(self):
        if os.path.exists(self.outputDir):
            raise ValueError("Output dir " + self.outputDir + " must not exist")

        os.makedirs(self.outputDir)

    def readYtypItems(self):
        self.ytypItems = YtypParser.readYtypDirectory(os.path.join(os.path.dirname(__file__), "..", "..", "resources", "ytyp"))
        self.lowercaseYtypItems = dict((k.lower(), k) for k, v in self.ytypItems.items())

    def repl(self, match: Match, fixedArchetypeNames: set[str]) -> str:
        archetypeName = match.group(2).lower()

        if archetypeName.lower() in self.lowercaseYtypItems and archetypeName not in self.ytypItems:
            fixedArchetypeName = self.lowercaseYtypItems[archetypeName.lower()]
            fixedArchetypeNames.add("changing archetypeName from " + archetypeName + " to " + fixedArchetypeName)
        else:
            fixedArchetypeName = archetypeName

        flags = int(match.group(4))
        origQuat = [float(match.group(9)), -float(match.group(6)), -float(match.group(7)), -float(match.group(8))]

        rotationQuaternion = np.divide(origQuat, [transforms3d.quaternions.qnorm(origQuat)])

        axangle = transforms3d.quaternions.quat2axangle(rotationQuaternion)
        rotationQuaternion = transforms3d.quaternions.axangle2quat(axangle[0], axangle[1])

        if np.allclose(rotationQuaternion, origQuat, rtol=0, atol=1e-05):
            rotationQuaternion = origQuat
        else:
            print("\t\tfixed rotation:", origQuat, rotationQuaternion)

        if transforms3d.quaternions.nearly_equivalent(rotationQuaternion, Sanitizer.identityQuaternion, rtol=0, atol=1e-05):
            rotationQuaternion = Sanitizer.identityQuaternion
        # do not remove this flag. it is required for dynamic objects as well (even if initially there is no custom orientation)
        # flags &= ~FLAG_ALLOW_FULL_ROTATION
        else:
            # TODO when is it necessary to add this flag? looking at some original rockstar maps only some rotations need this flag
            flags |= Flag.ALLOW_FULL_ROTATION

        return match.group(1) + \
               fixedArchetypeName + \
               match.group(3) + \
               str(flags) + \
               match.group(5) + \
               'x="' + Util.floatToStr(-rotationQuaternion[1]) + \
               '" y="' + Util.floatToStr(-rotationQuaternion[2]) + \
               '" z="' + Util.floatToStr(-rotationQuaternion[3]) + \
               '" w="' + Util.floatToStr(rotationQuaternion[0]) + '"' + \
               match.group(10) + Util.floatToStr(0) + match.group(11)

    def processFiles(self):
        for filename in natsorted(os.listdir(self.inputDir)):
            if filename.endswith(".ymap.xml"):
                self.processFile(filename)

    def processFile(self, filename: str):
        print("\tprocessing " + filename)

        f = open(os.path.join(self.inputDir, filename), 'r')
        content = f.read()
        f.close()

        fixedArchetypeNames = set()
        content_new = re.sub('(<Item type="CEntityDef">' +
                             '\\s*<archetypeName>)([^<]+)(</archetypeName>' +
                             '\\s*<flags value=")([^"]+)("\\s*/>' +
                             '(?:\\s*<[^/].*>)*?' +
                             '\\s*<rotation )x="([^"]+)" y="([^"]+)" z="([^"]+)" w="([^"]+)"(/>' +
                             '(?:\\s*<[^/].*>)*?' +
                             '\\s*<childLodDist value=")[^"]+("/>' +
                             '(?:\\s*<[^/].*>)*?' +
                             '\\s*</Item>)', lambda match: self.repl(match, fixedArchetypeNames), content, flags=re.M)

        for fixed in natsorted(fixedArchetypeNames):
            print("\t\t" + fixed)

        content_new = Ymap.replaceName(content_new, filename.lower()[:-9])
        content_new = Ymap.calculateAndReplaceLodDistance(content_new, self.ytypItems)
        content_new = Ymap.fixMapExtents(content_new, self.ytypItems)

        f = open(os.path.join(self.outputDir, filename.lower()), 'w')
        f.write(content_new)
        f.close()

    def copyOthers(self):
        # copy other files
        Util.copyFiles(self.inputDir, self.outputDir, lambda filename: not filename.endswith(".ymap.xml"))


from sklearn.metrics import classification_report
import numpy as np
import pandas as pd
from collections import defaultdict
from scipy.special import softmax

def metrics(p, return_df=False, plus_recall=0, tokenizer=None, dataset=None, merged_negative={0, 1, -100}, padding_label=-100, csize=15, subword_label=1,
            verbose=False):
    r''' TODO: This could be done better, for sure. But it works.
    '''
    predictions = np.array(p.predictions)
    predictions = softmax(predictions, axis=2)
    examples = None
    if plus_recall > 0:
        # Devalue 0 and 1
        predictions[:, :, 0] = predictions[:, :, 0] - (predictions[:, :, 0] * plus_recall)
        predictions[:, :, 1] = predictions[:, :, 1] - (predictions[:, :, 1] * plus_recall)

    preds = np.argmax(predictions, axis=2)
    # Ignore predictions where label == -100, padding
    preds[np.where(p.label_ids == -100)] = -100

    if dataset is not None and tokenizer is not None:
        examples = {'fp': defaultdict(list), 'fn': defaultdict(list)}
        ilbl = {v:k for k,v in tokenizer.label_map.items()}
        for i in range(len(preds)):
            st = None
            for j in range(len(preds[i])):
                _p = preds[i][j]
                _l = p.label_ids[i][j]
                if len(p.label_ids[i]) > (j+1) and p.label_ids[i][j+1] != padding_label:
                    _p2 = preds[i][j+1]
                    _l2 = p.label_ids[i][j+1]
                else:
                    _p2 = None
                    _l2 = None

                _d = dataset[i]['input_ids']
                id = dataset[i]['id']
                name = dataset[i]['name']
                if _l not in {subword_label, padding_label}: # We ignore padding and subwords
                    if _l != _p:
                        if st is None:
                            st = max(0, j-csize)
                            _j = j

                        if not (_l2 is not None and _l2 == _l and _l2 != _p2 and _p2 == _p):
                            # We want to merge tokens if it is the same label and same prediction when recording the examples, that is why we have the if
                            t = tokenizer.hf_tokenizer.decode(_d[st:_j]) + "<<" + str(tokenizer.hf_tokenizer.decode(_d[_j:j+1])) + \
                                ">>" + tokenizer.hf_tokenizer.decode(_d[j+1:j+csize])
                            value = str(tokenizer.hf_tokenizer.decode(_d[_j:j+1])).strip()
                            examples['fp'][ilbl[_p]].append(({'id': id, 'name': name, 'value': value,
                                                              'label': tokenizer.cui2name.get(ilbl[_l], ilbl[_l]), 'text': t}))
                            examples['fn'][ilbl[_l]].append(({'id': id, 'name': name, 'value': value,
                                                              'prediction': tokenizer.cui2name.get(ilbl[_p], ilbl[_p]), 'text': t}))

                            st = None

    _labels = np.reshape(p.label_ids, -1)
    _preds = np.reshape(preds, -1)

    cr = classification_report(_labels, _preds, output_dict=True)
    _cr = {}
    ignore = [str(tokenizer.label_map['O']), str(tokenizer.label_map['X']), '-100']
    ilabel_map = {str(v):k for k,v in tokenizer.label_map.items()}
    for key in cr.keys():
        if key not in ignore and key in ilabel_map:
            _cr[key] = cr[key]

    # Get merged metrics, ie all PID is just one entity
    tp = defaultdict(int)
    fn = defaultdict(int)
    tp_all = 0
    fp_all = 0
    for i, _label in enumerate(_labels):
        _pred = _preds[i]
        if _label in merged_negative:
            if _pred in merged_negative:
                tp[_label] += 1
            else:
                fn[_label] += 1

            if _label == 0:
                if _pred not in merged_negative:
                    fp_all += 1
        else:
            if _pred not in merged_negative:
                tp[_label] += 1
                tp_all += 1
            else:
                fn[_label] += 1
    for key in _cr:
        key = int(key)
        if int(key) in tp:
            _cr[str(key)]['r_merged'] = tp[key] / (tp[key] + fn.get(key, 0)) if tp[key] + fn.get(key, 0) > 0 else 0
        else:
            _cr[str(key)]['r_merged'] = None

    data = [['cui', 'name', 'p', 'r', 'f1', 'support', 'r_merged', 'p_merged']]
    for key in _cr:
        cui = ilabel_map[key]
        p_merged = tp_all / (tp_all + fp_all) if (tp_all + fp_all) > 0 else 0
        data.append([cui, tokenizer.cui2name.get(cui, cui), _cr[key]['precision'],
                     _cr[key]['recall'], _cr[key]['f1-score'], _cr[key]['support'], _cr[key]['r_merged'], p_merged])

    df = pd.DataFrame(data[1:], columns=data[0])
    if verbose:
        print(df)

    if not return_df:
        return {'recall': np.average(df.r.values), 'precision': np.average(df.p.values), 'f1': np.average(df.f1.values),
                'recall_merged': np.average([x for x in df.r_merged.values if pd.notna(x)]),
                'precison_merged': np.average([x for x in df.p_merged.values if pd.notna(x)])}
    else:
        return df, examples


import numpy as np
import cv2
import time
import facialemotions as emo

cap = cv2.VideoCapture(0)

while True:
   ret, frame = cap.read()
   if not ret:
      break
   frame = emo.extract_face(frame)
   #emo.extract_emotion(frame)
   cv2.imshow('frame', frame)
   cv2.waitKey(100)
   #cv2.waitKey(3000)

cap.release()
cv2.destroyAllWindows()


"""
This file contains the definition of different heterogeneous datasets used for training
"""
import torch
import numpy as np
import joblib

from .base_dataset import BaseDataset

class MixedDataset(torch.utils.data.Dataset):
    def __init__(self, options, dataset=None, **kwargs):
        self.dataset_list = ['h36m', 'lsp-orig', 'mpii', 'coco', 'youtube', 'mpi-inf-3dhp']
        self.datasets = [BaseDataset(options, ds, **kwargs) for ds in self.dataset_list]

        total_length = sum([len(ds) for ds in self.datasets])
        for ds_name, ds in zip(self.dataset_list, self.datasets):
            print("{} Train Num: {}".format(ds_name, format(len(ds), ",")))
        print("Total Train Num: {}".format(format(total_length, ",")))
        length_itw = sum([len(ds) for ds in self.datasets[1:-1]])
        self.length = max([len(ds) for ds in self.datasets])

        # self.amass = joblib.load("data/vibe_db/amass_db.pt")
        # self.amass = self.amass['theta']
        """
        Data distribution inside each batch:
        30% H36M - 60% ITW - 10% MPI-INF
        """
        if dataset == "pw_3d":
            self.partition = [.3, .6*len(self.datasets[1])/length_itw,
                            .6*len(self.datasets[2])/length_itw,
                            .6*len(self.datasets[3])/length_itw,
                            .6*len(self.datasets[4])/length_itw,
                            .6*len(self.datasets[5])/length_itw,
                            0.1]

        elif dataset == "h36m":
            self.partition = [1]

        elif dataset == "no_youtube":
            self.partition = [.3, .6*len(self.datasets[1])/length_itw,
                            .6*len(self.datasets[2])/length_itw,
                            .6*len(self.datasets[3])/length_itw,
                            0.1]
        else:
            self.partition = [.3, .6*len(self.datasets[1])/length_itw,
                            .6*len(self.datasets[2])/length_itw,
                            .6*len(self.datasets[3])/length_itw,
                            .6*len(self.datasets[4])/length_itw,
                            0.1]

        self.partition = np.array(self.partition).cumsum()

    def __getitem__(self, index):
        p = np.random.rand()
        for i in range(len(self.partition)):
            if p <= self.partition[i]:
                return self.datasets[i][index % len(self.datasets[i])]

    def __len__(self):
        return self.length


import face_recognition
import cv2
import os
import argparse
import face_recognition
import numpy as np
import demo_texture
import tensorflow as tf

from face_detection import select_face
from face_swap import face_swap

from api import PRN
from utils.render import render_texture

import numpy as np
import os
from glob import glob
import scipy.io as sio
from skimage.io import imread, imsave
from skimage.transform import rescale, resize
from time import time
import argparse
import ast
import matplotlib.pyplot as plt
import argparse

from api import PRN
from utils.render import render_texture
#from pytorch_stylegan_encoder import encode_image as ei
import cv2

def best_value(face_encodings, known_face_encodings):
    A = []
    for face_encoding in face_encodings:
        face_distances = face_recognition.face_distance(known_face_encodings, face_encoding)
        A.append(face_distances)
    return min(A)

def texture_editing(prn, image_path, ref_path, output_path, mode = 1):
    # read image
    image = imread(image_path)
    [h, w, _] = image.shape
    if h > 2000 or w > 2000:
        h = int(h*0.6)
        w = int(w*0.6)
        image = cv2.resize(image ,(w,h), interpolation=cv2.INTER_AREA)

    #-- 1. 3d reconstruction -> get texture.
    pos = prn.process(image)
    vertices = prn.get_vertices(pos)
    image = image/255.
    texture = cv2.remap(image, pos[:,:,:2].astype(np.float32), None, interpolation=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT,borderValue=(0))

    #-- 2. Texture Editing
    Mode = mode
    # change part of texture(for data augumentation/selfie editing. Here modify eyes for example)

    # change whole face(face swap)
    if Mode == 1:
        # texture from another image or a processed texture
#        ref_image = imread(ref_path)
        ref_image = ref_path
        ref_pos = prn.process(ref_image)
        ref_image = ref_image/255.
        ref_texture = cv2.remap(ref_image, ref_pos[:,:,:2].astype(np.float32), None, interpolation=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT,borderValue=(0))
        ref_vertices = prn.get_vertices(ref_pos)
        new_texture = ref_texture#(texture + ref_texture)/2.

    else:
        print('Wrong Mode! Mode should be 0 or 1.')
        exit()

    #-- 3. remap to input image.(render)
    vis_colors = np.ones((vertices.shape[0], 1))
    face_mask = render_texture(vertices.T, vis_colors.T, prn.triangles.T, h, w, c = 1)
    face_mask = np.squeeze(face_mask > 0).astype(np.float32)

    new_colors = prn.get_colors_from_texture(new_texture)
    new_image = render_texture(vertices.T, new_colors.T, prn.triangles.T, h, w, c = 3)
    new_image = image*(1 - face_mask[:,:,np.newaxis]) + new_image*face_mask[:,:,np.newaxis]

    # Possion Editing for blending image
    vis_ind = np.argwhere(face_mask>0)
    vis_min = np.min(vis_ind, 0)
    vis_max = np.max(vis_ind, 0)
    center = (int((vis_min[1] + vis_max[1])/2+0.5), int((vis_min[0] + vis_max[0])/2+0.5))
    output = cv2.seamlessClone((new_image*255).astype(np.uint8), (image*255).astype(np.uint8), (face_mask*255).astype(np.uint8), center, cv2.NORMAL_CLONE)

    # save output
    imsave(output_path, output)
    print('Done.')

minju_image = face_recognition.load_image_file("./known/jimin.jpg")
minju_face_encoding = face_recognition.face_encodings(minju_image)[0]

known_face_encodings = [
    minju_face_encoding,
]
known_face_names = [
    "Min-Ju"
]

# Load an image with an unknown face
unknown_image = face_recognition.load_image_file("./unknown/unknown.jpg")
#cnn 개느림
face_locations = face_recognition.face_locations(unknown_image)

face_encodings = face_recognition.face_encodings(unknown_image, face_locations)

best = best_value(face_encodings, known_face_encodings)

# 최고 닮은 도가 임계 값을 넘지 못하면 ㅂ2
if best > 0.5:
    best = 0

find_my = 0

for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
    matches = face_recognition.compare_faces(known_face_encodings, face_encoding)
    name = "Unknown"
    face_distances = face_recognition.face_distance(known_face_encodings, face_encoding)
    print(face_distances)
    #하나만 찾는데에는 쓸데없음
    best_match_index = np.argmin(face_distances)
    if face_distances == best:
        name = known_face_names[best_match_index]
        find_my = (top, right, bottom, left)
#445 1619 534 1529
#unknown_image
print(best_value(face_encodings, known_face_encodings))
print(find_my)

image = face_recognition.load_image_file("./unknown/unknown.jpg")
face_locations = face_recognition.face_locations(image)
print(face_locations)
img = cv2.imread("./unknown/unknown.jpg")
img2 = cv2.imread("./unknown/unknown.jpg")
print(face_locations[0][0])
print(face_locations[0][1])
print(face_locations[0][2])
print(face_locations[0][3])

"""
cv2.imshow("asdf", img[0:500, 100: 400])

cv2.waitKey(0)

cv2.imshow("asdf", img[face_locations[0][0]:face_locations[0][2],face_locations[0][3]:face_locations[0][1]])

cv2.waitKey(0)

cv2.destroyAllWindows()
"""

#! /usr/bin/env python
count_name = 0

for i in range(len(face_locations)):

    file_list = os.listdir("./data/")
    print(file_list)

    name_list = []
    for name in file_list:
        name_list.append(name)

    args_src = "./data/" + name_list[count_name]
    args_dst = "./unknown/unknown.jpg"
    args_out = "./output/result"

    print("count_name",count_name)
    print("face_locations",len(face_locations))
    print("face_locations[count_name]",len(face_locations[count_name]))

    Y = int((face_locations[count_name][2]-face_locations[count_name][0])/2)
    X = int((face_locations[count_name][1]-face_locations[count_name][3])/2)

    # Read images
    src_img = cv2.imread(args_src)
    #dst_img = cv2.imread(args.dst)

    src_img =cv2.cvtColor(src_img,cv2.COLOR_BGR2RGB)
    dst_img = img[face_locations[count_name][0]-Y:face_locations[count_name][2]+Y,face_locations[count_name][3]-X:face_locations[count_name][1]+X]
    # Select src face
    dst_img =cv2.cvtColor(dst_img,cv2.COLOR_BGR2RGB)
    os.environ['CUDA_VISIBLE_DEVICES'] = "0" # GPU number, -1 for CPU
    tf.reset_default_graph ()
    prn = PRN(is_dlib = True)

    texture_editing(prn, "./unknown/unknown.jpg", src_img, "./unknown/unknown.jpg", mode = 1)

    img =cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
    cv2.imwrite("asdf.jpg" , img)
    count_name += 1

"""

    if src_points is None or dst_points is None:
        print('Detect 0 Face !!!')
        exit(-1)

"""

img[find_my[0]:find_my[2],find_my[3]:find_my[1]] = img2[find_my[0]:find_my[2],find_my[3]:find_my[1]]

cv2.imwrite("AAAAAB.jpg" , img)
cv2.destroyAllWindows()