#!/usr/bin/env python

import argparse
from collections import defaultdict
import datetime
import math
import numpy as np
import os
import pandas as pd
import pickle
import pysam
import re
import time
import sys

from light_utils import *

def max_base(seq):
    base_counts = {"A": [], "T": [], "G": [], "C": []}
    for i in range(len(seq) - 9):
        for b in base_counts.keys():
            base_counts[b].append(seq[i : i + 10].count(b))
    return (
        max(base_counts["A"]),
        max(base_counts["T"]),
        max(base_counts["G"]),
        max(base_counts["C"]),
    )

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--bams",
        nargs="+",
        required=True,
        help="bams to parse (either one or two for paired end)",
    )
    parser.add_argument("--outpath", help="folder to write output to")
    parser.add_argument(
        "--UMI_bar", action="store_true", help="extract UMI and barcode"
    )
    parser.add_argument(
        "--paired",
        action="store_true",
        help="run once with each read primary and concatenate the files",
    )
    parser.add_argument(
        "--annotator", required=True, help="the path to the annotator pickle file"
    )
    parser.add_argument("--gtf", required=True, help="the path to the gtf file")
    parser.add_argument(
        "--stranded_library",
        action="store_true",
        help="Take gene strand information directly from the read strand",
    )

    args = parser.parse_args()
    return args

def extract_info_align(
    CI_dict,
    bam_read,
    suffix,
    bam_file,
    ann,
    UMI_bar,
    stranded_library,
    fill_char=np.nan,
):
    sec_dict = {True: 0, False: 1}
    if UMI_bar:
        vals = bam_read.query_name.split("_")
        CI_dict["barcode"].append(vals[-2])
        CI_dict["UMI"].append(vals[-1])
    else:
        CI_dict["barcode"].append(fill_char)
        CI_dict["UMI"].append(fill_char)
    CI_dict["id"].append(bam_read.query_name)
    CI_dict["fileType" + suffix].append("Aligned")
    CI_dict["readLen" + suffix].append(bam_read.query_length)
    CI_dict["aScore{}A".format(suffix)].append(bam_read.get_tag("AS"))
    CI_dict["NH{}A".format(suffix)].append(bam_read.get_tag("NH"))
    CI_dict["HI{}A".format(suffix)].append(bam_read.get_tag("HI"))
    CI_dict["nmm{}A".format(suffix)].append(nmm(bam_read.get_tag("MD")))
    CI_dict["qual{}A".format(suffix)].append(bam_read.mapping_quality)
    cigar1, cigar2 = split_cigar_align(bam_read.cigarstring)
    M, S, I, D = get_SM(cigar1)
    CI_dict["M{}A".format(suffix)].append(M)
    CI_dict["S{}A".format(suffix)].append(S)
    M, S, I, D = get_SM(cigar2)
    CI_dict["M{}B".format(suffix)].append(M)
    CI_dict["S{}B".format(suffix)].append(S)
    CI_dict["cigar{}A".format(suffix)].append(cigar1)
    CI_dict["cigar{}B".format(suffix)].append(cigar2)
    seq = bam_read.query_sequence
    counts = count_stretch(seq)
    CI_dict["AT_run_" + suffix].append(max(counts["A"], counts["T"]))
    CI_dict["GC_run_" + suffix].append(max(counts["G"], counts["C"]))
    CI_dict["max_run_" + suffix].append(max(counts.values()))
    CI_dict["seq{}".format(suffix)].append(seq)
    CI_dict["entropy" + suffix].append(float("{:.3f}".format(entropy(seq))))
    maxA, maxT, maxG, maxC = max_base(seq)
    CI_dict["maxA_10mer" + suffix].append(maxA)
    CI_dict["maxT_10mer" + suffix].append(maxT)
    CI_dict["maxG_10mer" + suffix].append(maxG)
    CI_dict["maxC_10mer" + suffix].append(maxC)

    refName, chrA, geneA, posA, chrB, geneB, posB = readObj_refname(
        bam_read.flag,
        bam_read.cigarstring,
        bam_file.get_reference_name(bam_read.tid),
        bam_read.reference_start + 1,
        ann,
        fill_char,
        stranded_library,
    )
    #   print("refName",refName)
    CI_dict["refName_AB" + suffix].append(refName)
    CI_dict["chr{}A".format(suffix)].append(chrA)
    CI_dict["chr{}B".format(suffix)].append(chrB)
    CI_dict["gene{}A".format(suffix)].append(geneA)
    CI_dict["gene{}B".format(suffix)].append(geneB)
    CI_dict["juncPos{}A".format(suffix)].append(int(posA))
    if np.isnan(posB):
        CI_dict["juncPos{}B".format(suffix)].append(posB)
    else:
        CI_dict["juncPos{}B".format(suffix)].append(int(posB))
    CI_dict["read_strand{}A".format(suffix)].append(read_strand(bam_read.flag))
    CI_dict["read_strand{}B".format(suffix)].append(fill_char)
    CI_dict["flag{}A".format(suffix)].append(bam_read.flag)
    CI_dict["flag{}B".format(suffix)].append(fill_char)

    CI_dict["primary{}A".format(suffix)].append(sec_dict[bam_read.is_secondary])
    CI_dict["primary{}B".format(suffix)].append(fill_char)
    empty_cols = [
        "aScore{}B".format(suffix),
        "qual{}B".format(suffix),
        "NH{}B".format(suffix),
        "nmm{}B".format(suffix),
        "HI{}B".format(suffix),
    ]
    for c in empty_cols:
        CI_dict[c].append(fill_char)
    return CI_dict

def extract_info_chim(
    CI_dict,
    bam_read1,
    bam_read2,
    suffix,
    bam_file,
    ann,
    UMI_bar,
    stranded_library,
    fill_char=np.nan,
):
    assert bam_read1.query_name == bam_read2.query_name
    sec_dict = {True: 0, False: 1}
    if UMI_bar:
        vals = bam_read1.query_name.split("_")
        CI_dict["barcode"].append(vals[-2])
        CI_dict["UMI"].append(vals[-1])
    else:
        CI_dict["barcode"].append(fill_char)
        CI_dict["UMI"].append(fill_char)
    reads = [bam_read1, bam_read2]
    CI_dict["fileType" + suffix].append("Chimeric")
    halves = ["A", "B"]
    CI_dict["id"].append(bam_read1.query_name)

    CI_dict["readLen" + suffix].append(bam_read1.query_length)

    seq = bam_read1.query_sequence
    counts = count_stretch(seq)
    CI_dict["AT_run_" + suffix].append(max(counts["A"], counts["T"]))
    CI_dict["GC_run_" + suffix].append(max(counts["G"], counts["C"]))
    CI_dict["max_run_" + suffix].append(max(counts.values()))
    CI_dict["seq{}".format(suffix)].append(seq)
    CI_dict["entropy" + suffix].append(float("%3.f" % (entropy(seq))))
    maxA, maxT, maxG, maxC = max_base(seq)
    CI_dict["maxA_10mer" + suffix].append(maxA)
    CI_dict["maxT_10mer" + suffix].append(maxT)
    CI_dict["maxG_10mer" + suffix].append(maxG)
    CI_dict["maxC_10mer" + suffix].append(maxC)

    refName, chrA, geneA, posA, chrB, geneB, posB = chim_refName(
        [x.flag for x in reads],
        [x.cigarstring for x in reads],
        [x.reference_start + 1 for x in reads],
        [bam_file.get_reference_name(x.tid) for x in reads],
        ann,
        stranded_library,
    )
    CI_dict["refName_AB" + suffix].append(refName)
    #   split_ref = refName.split("|")
    CI_dict["chr{}A".format(suffix)].append(chrA)
    CI_dict["chr{}B".format(suffix)].append(chrB)
    CI_dict["gene{}A".format(suffix)].append(geneA)
    CI_dict["gene{}B".format(suffix)].append(geneB)
    CI_dict["juncPos{}A".format(suffix)].append(int(posA))
    CI_dict["juncPos{}B".format(suffix)].append(int(posB))
    for i in range(2):
        CI_dict["aScore{}{}".format(suffix, halves[i])].append(reads[i].get_tag("AS"))
        CI_dict["qual{}{}".format(suffix, halves[i])].append(reads[i].mapping_quality)
        CI_dict["NH{}{}".format(suffix, halves[i])].append(reads[i].get_tag("NH"))
        CI_dict["HI{}{}".format(suffix, halves[i])].append(reads[i].get_tag("HI"))
        CI_dict["nmm{}{}".format(suffix, halves[i])].append(nmm(reads[i].get_tag("MD")))
        cigar = split_cigar_chim(reads[i].cigarstring)
        M, S, I, D = get_SM(cigar)
        CI_dict["M{}{}".format(suffix, halves[i])].append(M)
        CI_dict["S{}{}".format(suffix, halves[i])].append(S)
        CI_dict["cigar{}{}".format(suffix, halves[i])].append(cigar)
        CI_dict["read_strand{}{}".format(suffix, halves[i])].append(
            read_strand(reads[i].flag)
        )
        CI_dict["flag{}{}".format(suffix, halves[i])].append(reads[i].flag)

        CI_dict["primary{}{}".format(suffix, halves[i])].append(
            sec_dict[reads[i].is_secondary]
        )
    return CI_dict

def get_final_df(bam_files, j, suffixes, ann, UMI_bar, t0, gtf, stranded_library):
    CI_dfs = []
    for i in range(len(bam_files)):
        if i == 1:
            read_ids = set(CI_dfs[0]["id"])
        else:
            read_ids = set()
        suffix = suffixes[i]
        col_bases = [
            "aScore",
            "M",
            "S",
            "nmm",
            "qual",
            "NH",
            "HI",
            "cigar",
            "juncPos",
            "gene",
            "chr",
            "read_strand",
            "primary",
            "flag",
        ]
        columns = [
            "id",
            "readLen" + suffix,
            "fileType" + suffix,
            "seq" + suffix,
            "AT_run_" + suffix,
            "GC_run_" + suffix,
            "max_run_" + suffix,
            "entropy" + suffix,
            "refName_AB" + suffix,
            "UMI",
            "barcode",
            "maxA_10mer" + suffix,
            "maxT_10mer" + suffix,
            "maxG_10mer" + suffix,
            "maxC_10mer" + suffix,
        ]
        for c in col_bases:
            #     for r in ["R1"]:
            for l in ["A", "B"]:
                columns.append("{}{}{}".format(c, suffix, l))
        CI_dict = {c: [] for c in columns}
        count = 0
        first = False
        if i == 0:
            genomic_alignments = {}
        alignFile = pysam.AlignmentFile(bam_files[i])
        # columns
        for bam_read in alignFile.fetch(until_eof=True):
            #     suffix = "R1"

            # make sure read is mapped
            if not bam_read.is_unmapped:
                if (i == 0) or (
                    not bam_read.is_secondary and bam_read.query_name in read_ids
                ):
                    # it's a chimeric alignment and we need another line from it
                    if bam_read.has_tag("ch") and not first:
                        prev_read = bam_read
                        first = True
                    else:

                        # add info from chimeric read
                        if bam_read.has_tag("ch"):
                            count += 1

                            # note: removing chim for this test ONLY; uncomment after
                            CI_dict = extract_info_chim(
                                CI_dict,
                                prev_read,
                                bam_read,
                                suffix,
                                alignFile,
                                ann,
                                UMI_bar,
                                stranded_library,
                            )
                            first = False

                        # add info from align read
                        elif "N" in bam_read.cigarstring:
                            count += 1
                            CI_dict = extract_info_align(
                                CI_dict,
                                bam_read,
                                suffix,
                                alignFile,
                                ann,
                                UMI_bar,
                                stranded_library,
                            )

                        # save genomic alignment information
                        else:
                            if i == 0:
                                if bam_read.query_name not in genomic_alignments:
                                    genomic_alignments[
                                        bam_read.query_name
                                    ] = bam_read.get_tag("AS")
                                else:
                                    genomic_alignments[bam_read.query_name] = max(
                                        bam_read.get_tag("AS"),
                                        genomic_alignments[bam_read.query_name],
                                    )
                            else:
                                CI_dict = extract_info_align(
                                    CI_dict,
                                    bam_read,
                                    suffix,
                                    alignFile,
                                    ann,
                                    UMI_bar,
                                    stranded_library,
                                )
        #     if count == 10000:
        #       continue
        CI_df = pd.DataFrame.from_dict(CI_dict)
        if i == 0:
            genomic_alignments = defaultdict(lambda: np.nan, genomic_alignments)

            CI_df["genomicAlignmentR1"] = 0
            # print(CI_df.loc[CI_df["id"].isin(genomic_alignments)].shape)
            CI_df.loc[CI_df["id"].isin(genomic_alignments), "genomicAlignmentR1"] = 1
            CI_df["genomic_aScoreR1"] = CI_df["id"].map(genomic_alignments)
            CI_df["spliceDist"] = abs(CI_df["juncPosR1A"] - CI_df["juncPosR1B"])
        CI_dfs.append(CI_df)
    if len(bam_files) == 2:
        final_df = pd.merge(
            left=CI_dfs[0],
            right=CI_dfs[1][
                [c for c in CI_dfs[1].columns if c not in ["UMI", "barcode", "seqR2"]]
            ],
            how="left",
            left_on="id",
            right_on="id",
        )
        final_df["read_strand_compatible"] = 1
        final_df.loc[
            final_df["read_strandR1A"] == final_df["read_strandR2A"],
            "read_strand_compatible",
        ] = 0
        final_df["location_compatible"] = final_df.apply(get_loc_flag, axis=1)
    else:
        final_df = CI_dfs[0]
    #  final_df.fillna(np.nan,inplace=True)
    float_cols = [
        "aScoreR1A",
        "nmmR1A",
        "qualR1A",
        "NHR1A",
        "primaryR1A",
        "genomic_aScoreR1",
        "HIR1A",
    ]
    if len(bam_files) == 2:
        float_cols += [
            "readLenR2",
            "AT_run_R2",
            "GC_run_R2",
            "max_run_R2",
            "aScoreR2A",
            "aScoreR2B",
            "MR2A",
            "MR2B",
            "SR2A",
            "SR2B",
            "nmmR2A",
            "nmmR2B",
            "qualR2A",
            "qualR2B",
            "NHR2A",
            "NHR2B",
            "juncPosR2A",
            "juncPosR2B",
            "primaryR2A",
            "primaryR2B",
            "HIR2A",
            "HIR2B",
        ]

    #  for c in float_cols:
    #    final_df[c] = final_df[c].astype("Int32")
    print("started modify", time.time() - t0)
    #  for c in final_df.columns:
    #    str_dtype = final_df[c].dtype
    #    if str(str_dtype)[0] == "i":
    #      final_df[c] = final_df[c].astype("I" + str_dtype[1:])

    final_df = modify_refnames(final_df, gtf, stranded_library)

    print("ended modify", time.time() - t0)
    final_df["max_id_priority"] = final_df["id"].map(
        final_df.groupby("id")["HIR1A"].min()
    )

    #  for c in final_df.columns:
    #    if str(final_df[c].dtype)[0] == "I":
    #      final_df[c] = final_df[c].astype("float32")

    #  final_df.to_parquet(args.outpath + "class_input_final.pq")
    final_df["primary_bam"] = j
    final_df = final_df[final_df["HIR1A"] == final_df["max_id_priority"]]
    #  secondary = final_df[final_df["HIR1A"] != final_df["max_id_priority"]]

    return final_df

def main():
    t0 = time.time()
    args = get_args()

    bam_files = args.bams
    gtf = args.gtf
    print("bam_files", bam_files)

    annotator_path = args.annotator
    ann = pickle.load(open(annotator_path, "rb"))

    suffixes = ["R1", "R2"]

    final_dfs = []
    #  final_dfs_secondary = []

    if args.paired:
        n_rounds = 2
    else:
        n_rounds = 1

    for j in range(n_rounds):
        if j == 1:
            bam_files.reverse()
        print("bam_files", bam_files)
        primary = get_final_df(
            bam_files, j, suffixes, ann, args.UMI_bar, t0, gtf, args.stranded_library
        )
        final_dfs.append(primary)
    #    final_dfs_secondary.append(secondary)
    #  final_df = pd.concat(final_dfs,axis=0).reset_index(drop=True)

    #  for c in final_df.columns:
    #    if str(final_df[c].dtype)[0] == "I":
    #      final_df[c] = final_df[c].astype("float32")

    #  final_df.to_parquet(args.outpath + "class_input_final.pq")

    pd.concat(final_dfs, axis=0).reset_index(drop=True).to_parquet(
        args.outpath + "class_input.pq"
    )
    #  pd.concat(final_dfs_secondary,axis=0).reset_index(drop=True).to_parquet(args.outpath + "class_input_secondary.pq")

    #  final_df[final_df["HIR1A"] == final_df["max_id_priority"]].to_hdf(args.outpath + "class_input.h5", key="class_input")
    #  final_df[final_df["HIR1A"] != final_df["max_id_priority"]].to_hdf(args.outpath + "class_input_secondary.h5", key="class_input_secondary")

    pd.concat(final_dfs, axis=0).reset_index(drop=True).to_csv(
        args.outpath + "class_input.tsv", sep="\t", index=False
    )
    #  pd.concat(final_dfs_secondary,axis=0).reset_index(drop=True).to_csv(args.outpath + "class_input_secondary.tsv",sep="\t",index=False)
    #  final_df["juncPosR1A"] = final_df["juncPosR1A"].astype("Int32")
    #  final_df["juncPosR1B"] = final_df["juncPosR1B"].astype("Int32")

    print("total time", datetime.timedelta(seconds=time.time() - t0))

main()


##############################################################################
#
#  File format versions:
#
#  1: initial version
#
#  2: now contains reciprocal planck opacity and rosseland opacity
#     previous rosseland opacity was actually reciprocal planck opacity
#
##############################################################################

from __future__ import print_function, division

import os
import hashlib

import h5py
import numpy as np
from astropy import log as logger
import six

from ..version import __version__

from ..util.constants import c
from ..util.functions import FreezableClass
from ..util.interpolate import interp1d_fast_loglog
from ..util.integrate import integrate_loglog
from ..util.nans import check_for_nans

from .optical_properties import OpticalProperties
from .emissivities import Emissivities
from .mean_opacities import MeanOpacities

def henyey_greenstein(mu, g, p_lin_max):
    P1 = (1. - g * g) / (1. + g * g - 2. * g * mu) ** 1.5
    P2 = - p_lin_max * P1 * (1. - mu * mu) / (1. + mu * mu)
    P3 = P1 * 2. * mu / (1. + mu * mu)
    P4 = 0.
    return P1, P2, P3, P4

class SphericalDust(FreezableClass):
    r"""
    This class should be used in cases where fully arbitrary dust properties
    need to be specified, within the framework of randomly oriented grains,
    which means that the scattering phase function has the general form:

    .. math:: R(\theta) = \left[\begin{array}{cccc}P_1 & P_2 & 0 & 0 \\P_2 & P_1 & 0 & 0 \\0 & 0 & P_3 & -P_4 \\0 & 0 & P_4 & P_3\end{array}\right]

    This class is initialized with::

        d = SphericalDust()

    and the properties should then be set manually. See
    `here <http://docs.hyperion-rt.org/en/stable/setup/setup_dust.html#fully-customized-4-element-dust>`_
    for a description of the available properties and how to set them.
    """

    def __init__(self, *args):

        self._file = None
        self.md5 = None

        self.optical_properties = OpticalProperties()
        self.mean_opacities = MeanOpacities()
        self.emissivities = Emissivities()

        self.set_sublimation_specific_energy('no', 0.)

        self._freeze()

        if len(args) == 0:
            pass
        elif len(args) == 1:
            self.read(args[0])
        else:
            raise Exception("SphericalDust cannot take more than one argument")

    def hash(self):

        h = hashlib.md5()

        if self.optical_properties is None or not self.optical_properties.all_set():
            h.update('none'.encode('utf-8'))
        else:
            h.update(self.optical_properties.hash().encode('utf-8'))

        if self.emissivities is None or not self.emissivities.all_set():
            h.update('none'.encode('utf-8'))
        else:
            h.update(self.emissivities.hash().encode('utf-8'))

        if self.mean_opacities is None or not self.mean_opacities.all_set():
            h.update('none'.encode('utf-8'))
        else:
            h.update(self.mean_opacities.hash().encode('utf-8'))

        import struct
        h.update(self.sublimation_mode.encode('utf-8'))
        h.update(struct.pack('>d', self.sublimation_energy))

        return h.hexdigest()

    def set_lte_emissivities(self, n_temp=1200, temp_min=0.1, temp_max=100000.):
        '''
        Calculate the emissivities assuming LTE

        Parameters
        ----------
        n_temp : int, optional
            The number of temperatures to calculate the emissivities for
        temp_min : float, optional
            The minimum temperature to calculate the emissivities for
        temp_max : float, optional
            The maximum temperature to calculate the emissivities for
        '''
        self.mean_opacities.compute(self.optical_properties, n_temp=n_temp,
                                    temp_min=temp_min, temp_max=temp_max)
        self.emissivities.set_lte(self.optical_properties, self.mean_opacities)

    def plot(self, filename):

        # Check that the optical properties have been set
        self.optical_properties.ensure_all_set()

        import matplotlib.pyplot as plt

        # Save original rc parameters
        rc_orig = plt.rcParams

        # Reset to defaults
        plt.rcdefaults()
        plt.rc('legend', fontsize=7)
        plt.rc('axes', titlesize='x-small')
        plt.rc('axes', labelsize='x-small')
        plt.rc('xtick', labelsize='xx-small')
        plt.rc('ytick', labelsize='xx-small')
        plt.rc('axes', linewidth=0.5)
        plt.rc('patch', linewidth=0.5)

        # Compute mean opacities if not already existent
        self._compute_mean_opacities()

        # Check that emissivities are set (before computing mean opacities)
        if not self.emissivities.all_set():
            logger.info("Computing emissivities assuming LTE")
            self.emissivities.set_lte(self.optical_properties, self.mean_opacities)

        # Initialize figure
        fig = plt.figure(figsize=(10, 12))

        # Plot optical properties
        fig = self.optical_properties.plot(fig, [421, 423, 424, 425, 426])

        # Plot mean opacities
        fig = self.mean_opacities.plot(fig, 428)

        # Plot emissivities
        fig = self.emissivities.plot(fig, 427)

        # Adjust spacing between subplots
        fig.subplots_adjust(left=0.08, right=0.92, wspace=0.22, hspace=0.30)

        # Save figure
        fig.savefig(filename, bbox_inches='tight')

        # Close figure to save RAM
        plt.close(fig)

        # Restore rc parameters
        plt.rc(rc_orig)

    def set_sublimation_temperature(self, mode, temperature=0.):
        '''
        Set the dust sublimation mode and temperature.

        Parameters
        ----------
        mode : str
            The dust sublimation mode, which can be:
                * 'no'   - no sublimation
                * 'fast' - remove all dust in cells exceeding the
                           sublimation temperature
                * 'slow' - reduce the dust in cells exceeding the
                           sublimation temperature
                * 'cap'  - any temperature exceeding the sublimation
                           temperature is reset to the sublimation
                           temperature.

        temperature : float, optional
            The dust sublimation temperature, in K
        '''

        if mode not in ['no', 'fast', 'slow', 'cap']:
            raise Exception("mode should be one of no/fast/slow/cap")

        if mode != 'no' and temperature is None:
            raise Exception("Need to specify a sublimation temperature")

        self.sublimation_mode = mode
        self.sublimation_energy = self.temperature2specific_energy(temperature)

    def set_sublimation_specific_energy(self, mode, specific_energy=0.):
        '''
        Set the dust sublimation mode and specific energy.

        Parameters
        ----------
        mode : str
            The dust sublimation mode, which can be:
                * 'no'   - no sublimation
                * 'fast' - remove all dust in cells exceeding the
                           sublimation specific energy
                * 'slow' - reduce the dust in cells exceeding the
                           sublimation specific energy
                * 'cap'  - any specific energy exceeding the sublimation
                           specific energy is reset to the sublimation
                           specific energy.

        specific_energy : float, optional
            The dust sublimation specific energy, in cgs
        '''

        if mode not in ['no', 'fast', 'slow', 'cap']:
            raise Exception("mode should be one of no/fast/slow/cap")

        if mode != 'no' and specific_energy is None:
            raise Exception("Need to specify a sublimation specific_energy")

        self.sublimation_mode = mode
        self.sublimation_energy = specific_energy

    def _write_dust_sublimation(self, group):
        group.attrs['sublimation_mode'] = np.string_(self.sublimation_mode)
        if self.sublimation_mode in ['slow', 'fast', 'cap']:
            group.attrs['sublimation_specific_energy'] = self.sublimation_energy

    def _read_dust_sublimation(self, group):
        if 'sublimation_mode' in group.attrs:
            self.sublimation_mode = group.attrs['sublimation_mode'].decode('ascii')
            if self.sublimation_mode in ['slow', 'fast', 'cap']:
                self.sublimation_energy = group.attrs['sublimation_specific_energy']

    def _compute_mean_opacities(self):
        if not self.mean_opacities.all_set():
            self.mean_opacities.compute(self.optical_properties)

    def write(self, filename, compression=True):
        '''
        Write out to a standard dust file, including calculations of the mean
        opacities and optionally thermal emissivities.
        '''

        # Check that the optical properties have been set
        self.optical_properties.ensure_all_set()

        # Compute mean opacities if not already existent
        self._compute_mean_opacities()

        # Check that emissivities are set (before computing mean opacities)
        if not self.emissivities.all_set():
            logger.info("Computing emissivities assuming LTE")
            self.emissivities.set_lte(self.optical_properties, self.mean_opacities)

        # Create dust table set
        if isinstance(filename, six.string_types):
            dt = h5py.File(filename, 'w')
        else:
            dt = filename

        # Add standard keywords to header
        dt.attrs['version'] = 2
        dt.attrs['type'] = 1
        dt.attrs['python_version'] = np.string_(__version__)
        if self.md5:
            dt.attrs['asciimd5'] = np.string_(self.md5)

        # Add optical properties and scattering angle tables
        self.optical_properties.to_hdf5_group(dt)

        # Add mean opacities table
        self.mean_opacities.to_hdf5_group(dt)

        # Add emissivities and emissivity variable tables
        self.emissivities.to_hdf5_group(dt)

        # Dust sublimation parameters
        self._write_dust_sublimation(dt)

        # Check that there are no NaN values in the file - if there are, a
        # warning is emitted.
        check_for_nans(dt)

        # Close dust file
        if isinstance(dt, h5py.File):
            dt.close()

        self._file = (filename, self.hash())

    def read(self, filename):
        '''
        Read in from a standard dust file
        '''

        from ..util.functions import asstr

        if isinstance(filename, six.string_types):

            # Check file exists
            if not os.path.exists(filename):
                raise Exception("File not found: %s" % filename)

            # Read in dust table set
            dt = h5py.File(filename, 'r')
            close = True

        else:

            # Read in dust table set
            dt = filename
            close = False

        # Check version and type
        if dt.attrs['version'] not in [1, 2]:
            raise Exception("Version should be 1 or 2")
        if dt.attrs['type'] != 1:
            raise Exception("Type should be 1")
        if 'asciimd5' in dt.attrs:
            self.md5 = asstr(dt.attrs['asciimd5'])
        else:
            self.md5 = None

        # Read in the optical properties
        self.optical_properties.from_hdf5_group(dt)

        # Read in the planck and rosseland mean opacities
        if dt.attrs['version'] == 1:
            logger.warning("Version 1 dust file detected - discarding mean opacities and recomputing them")
            self.mean_opacities.compute(self.optical_properties)
        else:
            self.mean_opacities.from_hdf5_group(dt)

        # Read in emissivities
        self.emissivities.from_hdf5_group(dt)

        # Dust sublimation parameters
        self._read_dust_sublimation(dt)

        # Close file object if needed
        if close:
            dt.close()
            self._file = (filename, self.hash())

    def chi_nu_temperature(self, temperature):
        """
        Compute the mean opacity to extinction for a blackbody at a given temperature.

        Parameters
        ----------
        temperature : float
            The temperature of the blackbody to use

        Returns
        -------
        chi_nu_mean : float
            The mean opacity to extinction
        """
        self._compute_mean_opacities()
        return interp1d_fast_loglog(self.mean_opacities.temperature,
                                    self.mean_opacities.chi_planck,
                                    temperature,
                                    bounds_error=True)

    def kappa_nu_temperature(self, temperature):
        """
        Compute the mean opacity to absorption for a blackbody at a given temperature.

        Parameters
        ----------
        temperature : float
            The temperature of the blackbody to use

        Returns
        -------
        kappa_nu_mean : float
            The mean opacity to absorption
        """
        self._compute_mean_opacities()
        return interp1d_fast_loglog(self.mean_opacities.temperature,
                                    self.mean_opacities.kappa_planck,
                                    temperature,
                                    bounds_error=True)

    def chi_nu_spectrum(self, nu, fnu):
        """
        Compute the mean opacity to extinction for a given spectrum.

        Parameters
        ----------
        nu : array_like
            The frequencies, in Hz
        fnu : array_like
            The monochromatic fluxes per unit frequency. Units are unimportant
            since proportionality constants are cancelled out in the
            computation.

        Returns
        -------
        chi_nu_mean : float
            The mean opacity to extinction
        """
        if nu.max() > self.optical_properties.nu.max() or nu.min() < self.optical_properties.nu.min():
            raise Exception("Opacity to extinction is not defined at all "
                            "spectrum frequencies")
        chi_nu = self.optical_properties.interp_chi_nu(nu)
        return (integrate_loglog(nu, fnu * chi_nu) /
                integrate_loglog(nu, fnu))

    def kappa_nu_spectrum(self, nu, fnu):
        """
        Compute the mean opacity to absorption for a given spectrum.

        Parameters
        ----------
        nu : array_like
            The frequencies, in Hz
        fnu : array_like
            The monochromatic fluxes per unit frequency. Units are unimportant
            since proportionality constants are cancelled out in the
            computation.

        Returns
        -------
        kappa_nu_mean : float
            The mean opacity to absorption
        """
        if nu.max() > self.optical_properties.nu.max() or nu.min() < self.optical_properties.nu.min():
            raise Exception("Opacity to absorption is not defined at all "
                            "spectrum frequencies")
        kappa_nu = self.optical_properties.interp_kappa_nu(nu)
        return (integrate_loglog(nu, fnu * kappa_nu) /
                integrate_loglog(nu, fnu))

    def temperature2specific_energy(self, temperature):
        """
        Convert a temperature to its corresponding specific energy value.

        Parameters
        ----------
        temperature : float or array_like
            The temperature to convert

        Returns
        -------
        specific_energy : float or array_like
            The specific energy corresponding to the input temperature
        """

        self._compute_mean_opacities()

        specific_energy = interp1d_fast_loglog(self.mean_opacities.temperature,
                                               self.mean_opacities.specific_energy,
                                               temperature,
                                               bounds_error=False,
                                               fill_value=np.nan)

        if np.isscalar(temperature):
            if temperature < self.mean_opacities.temperature[0]:
                specific_energy = self.mean_opacities.specific_energy[0]
            elif temperature > self.mean_opacities.temperature[-1]:
                specific_energy = self.mean_opacities.specific_energy[-1]
        else:
            specific_energy[temperature < self.mean_opacities.temperature[0]] = self.mean_opacities.specific_energy[0]
            specific_energy[temperature > self.mean_opacities.temperature[-1]] = self.mean_opacities.specific_energy[-1]

        return specific_energy

    def specific_energy2temperature(self, specific_energy):
        """
        Convert a specific energy value to its corresponding temperature.

        Parameters
        ----------
        specific_energy : float or array_like
            The specific energy to convert

        Returns
        -------
        temperature : float or array_like
            The temperature corresponding to the input specific energy
        """

        self._compute_mean_opacities()

        temperature = interp1d_fast_loglog(self.mean_opacities.specific_energy,
                                           self.mean_opacities.temperature,
                                           specific_energy,
                                           bounds_error=False,
                                           fill_value=np.nan)

        if np.isscalar(specific_energy):
            if specific_energy < self.mean_opacities.specific_energy[0]:
                temperature = self.mean_opacities.temperature[0]
            elif specific_energy > self.mean_opacities.specific_energy[-1]:
                temperature = self.mean_opacities.temperature[-1]
        else:
            temperature[specific_energy < self.mean_opacities.specific_energy[0]] = self.mean_opacities.temperature[0]
            temperature[specific_energy > self.mean_opacities.specific_energy[-1]] = self.mean_opacities.temperature[-1]

        return temperature

class IsotropicDust(SphericalDust):
    """
    This class should be used for dust properties that include isotropic
    scattering. The dust properties should be instatiated as::

        d = IsotropicDust(nu, albedo, chi)

    where ``nu``, ``albedo``, and ``chi`` are 1-D Numpy arrays containing the
    frequencies, albedo, and opacity to extinction respectively.
    """

    def __init__(self, nu, albedo, chi):

        SphericalDust.__init__(self)

        # Set cos(theta) grid for computing the scattering matrix elements
        self.optical_properties.mu = np.linspace(-1., 1., 2)

        # Set optical properties
        self.optical_properties.nu = nu
        self.optical_properties.albedo = albedo
        self.optical_properties.chi = chi

        # Compute scattering matrix elements
        self.optical_properties.initialize_scattering_matrix()

        # Set scattering matrix to isotropic values
        self.optical_properties.P1[:, :] = 1.
        self.optical_properties.P2[:, :] = 0.
        self.optical_properties.P3[:, :] = 1.
        self.optical_properties.P4[:, :] = 0.

        # Sort optical properties
        self.optical_properties._sort()

class HenyeyGreensteinDust(SphericalDust):
    """
    This class should be used for dust properties that include
    scattering parameterized by the `Henyey-Greenstein, 1941
    <http://dx.doi.org/10.1086/144246>`_ function. The dust properties should
    be instatiated as::

        d = HenyeyGreensteinDust(nu, albedo, chi, g, p_lin_max)

    where ``nu``, ``albedo``, and ``chi`` are 1-D Numpy arrays containing the
    frequencies, albedo, and opacity to extinction respectively, and ``g`` and
    ``p_lin_max`` are also 1-D Numpy arrays containing the asymmetry parameter
    and the maximum linear polarization.
    """

    def __init__(self, nu, albedo, chi, g, p_lin_max):

        SphericalDust.__init__(self)

        # Set cos(theta) grid for computing the scattering matrix elements
        n_mu = 100
        self.optical_properties.mu = np.linspace(-1., 1., n_mu)

        # Set optical properties
        self.optical_properties.nu = nu
        self.optical_properties.albedo = albedo
        self.optical_properties.chi = chi

        # Compute scattering matrix elements
        self.optical_properties.initialize_scattering_matrix()

        for i in range(n_mu):
            self.optical_properties.P1[:, i], \
                self.optical_properties.P2[:, i], \
                self.optical_properties.P3[:, i], \
                self.optical_properties.P4[:, i] = henyey_greenstein(self.optical_properties.mu[i], g, p_lin_max)

class HOCHUNKDust(HenyeyGreensteinDust):
    """
    This class should be used for dust properties that include
    scattering parameterized by the `Henyey-Greenstein, 1941
    <http://dx.doi.org/10.1086/144246>`_ function, which are formatted for the
    `HOCHUNK code <http://gemelli.colorado.edu/~bwhitney/codes/>`_. The dust
    properties should be instatiated as::

        d = HOCHUNKDust(filename)

    where ``filename`` is the name of the file containing the dust properties
    in the HOCHUNK format.
    """

    def __init__(self, filename):

        # Read in dust file
        dustfile = np.loadtxt(filename, dtype=[('wav', float), ('c_ext', float),
                                               ('c_sca', float), ('chi', float), ('g', float),
                                               ('p_lin_max', float)], usecols=[0, 1, 2, 3, 4, 5])

        # Ensure file is ordered in increasing frequency
        if dustfile['wav'][-1] > dustfile['wav'][0]:
            dustfile = dustfile[::-1]

        # Compute frequency and albedo
        nu = c / dustfile['wav'] * 1.e4
        albedo = dustfile['c_sca'] / dustfile['c_ext']

        self.md5 = hashlib.md5(open(filename, 'rb').read()).hexdigest()

        HenyeyGreensteinDust.__init__(self, nu, albedo, dustfile['chi'], dustfile['g'], dustfile['p_lin_max'])

TTsreDust = HOCHUNKDust

class CoatsphSingle(SphericalDust):

    def __init__(self, directory, size, density):
        '''
        Initialize single-component dust.

        Parameters
        ----------
        directory : str
            Directory containing all the files describing the dust
        size : float
            Grain size, in cm
        density : float
            Dust grain density, in g/cm^3
        '''

        SphericalDust.__init__(self)

        f = open('%s/coatsph_forw.dat' % directory, 'rb')
        version = f.readline()
        n_components = int(f.readline().strip().split()[5])

        # Read in main dust file

        dustfile = np.loadtxt(f, skiprows=3,
                              dtype=[('x', float), ('radius', float), ('wav', float),
                                     ('q_ext', float), ('q_sca', float), ('q_back', float),
                                     ('g', float)])

        n_wav = len(dustfile)

        self.optical_properties.nu = c / dustfile['wav'] * 1.e4
        self.optical_properties.albedo = dustfile['q_sca'] / dustfile['q_ext']
        self.optical_properties.chi = 0.75 * dustfile['q_ext'] / size / density

        # Read in scattering matrix elements

        for i in range(n_wav):

            filename = '%s/coatsph_scat_%04i_0001.dat' % (directory, i + 1)

            phasefile = np.loadtxt(filename, skiprows=9,
                                   dtype=[('theta', float), ('s11', float), ('polariz',
                                                                             float), ('s12', float), ('s33', float), ('s34',
                                                                                float)])

            if i == 0:
                self.optical_properties.mu = np.cos(np.radians(phasefile['theta']))
                self.optical_properties.initialize_scattering_matrix()

            self.optical_properties.P1[i, :] = phasefile['s11']
            self.optical_properties.P2[i, :] = phasefile['s12']
            self.optical_properties.P3[i, :] = phasefile['s33']
            self.optical_properties.P4[i, :] = phasefile['s34']

class CoatsphMultiple(SphericalDust):

    def __init__(self, directory):
        '''
        Initialize multi-component dust.

        Parameters
        ----------
        directory : str
            Directory containing all the files describing the dust
        '''

        SphericalDust.__init__(self)

        f = open('%s/coatsph_forw.dat' % directory, 'rb')
        version = f.readline()
        n_components = int(f.readline().strip().split()[5])

        # Read in main dust file

        dustfile = np.loadtxt(f, skiprows=7,
                              dtype=[('wav', float), ('c_ext', float), ('c_sca', float),
                                     ('chi', float), ('g', float), ('pmax', float),
                                     ('thetmax', float)])

        n_wav = len(dustfile)
        self.optical_properties.nu = c / dustfile['wav'] * 1.e4
        self.optical_properties.albedo = dustfile['c_sca'] / dustfile['c_ext']
        self.optical_properties.chi = dustfile['chi']

        # Read in scattering matrix elements

        for i in range(n_wav):

            filename = '%s/coatsph_scat.%04i.dat' % (directory, i + 1)

            phasefile = np.loadtxt(filename, skiprows=7,
                                   dtype=[('theta', float), ('s11', float), ('polariz',
                                                                             float), ('s12', float), ('s33', float), ('s34',
                                                                                float)])

            if i == 0:
                self.optical_properties.mu = np.cos(np.radians(phasefile['theta']))
                self.optical_properties.initialize_scattering_matrix()

            self.optical_properties.P1[i, :] = phasefile['s11']
            self.optical_properties.P2[i, :] = phasefile['s12']
            self.optical_properties.P3[i, :] = phasefile['s33']
            self.optical_properties.P4[i, :] = phasefile['s34']

class MieXDust(SphericalDust):

    def __init__(self, model):

        SphericalDust.__init__(self)

        wav = np.loadtxt('%s.alb' % model, usecols=[0])
        self.optical_properties.albedo = np.loadtxt('%s.alb' % model, usecols=[1])
        kappa = np.loadtxt('%s.k_abs' % model, usecols=[1])
        self.optical_properties.chi = kappa / (1 - self.optical_properties.albedo)

        # Check for NaN values
        for quantity in ['chi', 'albedo']:

            values = self.optical_properties.__dict__[quantity]

            if np.any(np.isnan(values)):
                logger.warning("NaN values found inside MieX %s file - interpolating" % quantity)
                invalid = np.isnan(values)
                values[invalid] = interp1d_fast_loglog(wav[~invalid], values[~invalid], wav[invalid])
                if np.any(np.isnan(values)):
                    raise Exception("Did not manage to fix NaN values in MieX %s" % quantity)

        self.optical_properties.nu = c / wav * 1.e4

        n_wav = len(wav)
        n_mu = (len(open('%s.f11' % model).readlines()) // n_wav) - 1

        mu = np.zeros(n_mu)

        # Read mu
        f11 = open('%s.f11' % model)
        f11.readline()
        f11.readline()
        for i in range(n_mu):
            mu[i] = np.cos(np.radians(float(f11.readline().split()[0])))
        f11.close()
        self.optical_properties.mu = mu[::-1]

        # Read in matrix elements

        self.optical_properties.initialize_scattering_matrix()

        f11 = open('%s.f11' % model)
        f12 = open('%s.f12' % model)
        f33 = open('%s.f33' % model)
        f34 = open('%s.f34' % model)

        f11.readline()
        f12.readline()
        f33.readline()
        f34.readline()

        for j in range(n_wav):

            if float(f11.readline()) != wav[j]:
                raise Exception("Incorrect wavelength in f11")
            if float(f12.readline()) != wav[j]:
                raise Exception("Incorrect wavelength in f12")
            if float(f33.readline()) != wav[j]:
                raise Exception("Incorrect wavelength in f33")
            if float(f34.readline()) != wav[j]:
                raise Exception("Incorrect wavelength in f34")

            for i in range(n_mu):

                self.optical_properties.P1[j, n_mu - i - 1] = float(f11.readline().split()[1])
                self.optical_properties.P2[j, n_mu - i - 1] = float(f12.readline().split()[1])
                self.optical_properties.P3[j, n_mu - i - 1] = float(f33.readline().split()[1])
                self.optical_properties.P4[j, n_mu - i - 1] = float(f34.readline().split()[1])

        for i in range(n_mu):

            for quantity in ['P1', 'P2', 'P3', 'P4']:

                values = self.optical_properties.__dict__[quantity]

                if np.any(np.isnan(values[:, i])):
                    logger.warning("NaN values found inside MieX %s file - interpolating" % quantity)
                    invalid = np.isnan(values[:, i])
                    values[:, i][invalid] = interp1d_fast_loglog(wav[~invalid], values[:, i][~invalid], wav[invalid])
                    if np.any(np.isnan(values[:, i])):
                        raise Exception("Did not manage to fix NaN values in MieX %s" % quantity)

class BHDust(SphericalDust):
    """
    This class should be used for dust properties that were computed using
    `this dust calculation code <https://github.com/hyperion-rt/bhmie>`_ which
    is a wrapper to the ``bhmie`` routine originally written by C.F. Bohren and
    D. Huffman and improved by B. Draine.

    When using the ``bhmie`` code, you should set the output format to ``2``,
    which will create a number of files ending in ``.wav``, ``.mu``, ``.alb``,
    etc. Then, instantiate this class with the name of the directory containing
    these output files along with the prefix used. For example, if you use
    ``directory/mydust`` as a prefix in ``bhmie``, you can import this dust
    with::

        >>> from hyperion.dust import BHDust
        >>> d = BHDust('directory/mydust')

    """

    def __init__(self, model):

        SphericalDust.__init__(self)

        mu = np.loadtxt('%s.mu' % model)

        nu = c / np.loadtxt('%s.wav' % model) * 1.e4
        albedo = np.loadtxt('%s.alb' % model)
        chi = np.loadtxt('%s.chi' % model)

        P1 = np.loadtxt('%s.f11' % model)
        P2 = np.loadtxt('%s.f12' % model)
        P3 = np.loadtxt('%s.f33' % model)
        P4 = np.loadtxt('%s.f34' % model)

        if nu[-1] < nu[0]:
            nu = nu[::-1]
            albedo = albedo[::-1]
            chi = chi[::-1]
            P1 = P1[::-1, :]
            P2 = P2[::-1, :]
            P3 = P3[::-1, :]
            P4 = P4[::-1, :]

        if mu[-1] < mu[0]:
            mu = mu[::-1]
            P1 = P1[:, ::-1]
            P2 = P2[:, ::-1]
            P3 = P3[:, ::-1]
            P4 = P4[:, ::-1]

        self.optical_properties.mu = mu

        self.optical_properties.nu = nu
        self.optical_properties.albedo = albedo
        self.optical_properties.chi = chi

        self.optical_properties.P1 = P1
        self.optical_properties.P2 = P2
        self.optical_properties.P3 = P3
        self.optical_properties.P4 = P4


import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.data.dataset import Dataset
import os.path as osp
from PIL import Image
import numpy as np
from torchvision import transforms
from torchvision import utils
from utils import pose_utils
from PIL import ImageDraw
from utils.transforms import create_part
import time
import json
import random
import cv2
from data.base_dataset import BaseDataset, get_transform
np.seterr(divide='ignore', invalid='ignore')

class DemoDataset(BaseDataset):
    def __init__(self, config, augment):
        self.opt = config
        self.transforms = augment
        self.isval = self.opt.isval
        self.isdemo = self.opt.isdemo
        self.train_mode = self.opt.train_mode
        self.fine_width = 192
        self.fine_height = 256
        self.size = (256, 192)

        if self.opt.warp_cloth:
            self.img_list =[i.strip() for i in open('dataset/data_pair.txt', 'r').readlines()]
        else:
            self.img_list =[i.strip() for i in open('demo/demo.txt', 'r').readlines()]

    def __getitem__(self, index):
        t0 = time.time()
        try:
            img_source = self.img_list[index].split('\t')[0]
            target_pose = self.img_list[index].split('\t')[1]
            cloth_img = self.img_list[index].split('\t')[2]
            data_suffix = self.img_list[index].split('\t')[3]
        except:
            img_source = self.img_list[index].split(' ')[0]
            target_pose = self.img_list[index].split(' ')[1]
            cloth_img = self.img_list[index].split(' ')[2]
            data_suffix = self.img_list[index].split(' ')[3]

        data_suffix = data_suffix if data_suffix == 'train' else 'val'
        source_splitext = os.path.splitext(img_source)[0]
        cloth_splitext = os.path.splitext(cloth_img)[0]

        if self.opt.warp_cloth:
            source_splitext = os.path.join(data_suffix, source_splitext)
            cloth_img = os.path.join(data_suffix, cloth_img)
            cloth_splitext = os.path.join(data_suffix, cloth_splitext)

        source_img_path = os.path.join('dataset/images', source_splitext + '.jpg')
        cloth_parse_path = os.path.join('dataset/cloth_mask', cloth_splitext + '_mask.png')

        if self.opt.warp_cloth:
            cloth_img_path = os.path.join('dataset/images', cloth_img)
        else:
            cloth_img_path = os.path.join('dataset/cloth_image', cloth_img)

        ### image
        source_img = self.open_transform(source_img_path, False)
        cloth_img = self.open_transform(cloth_img_path, False)
        cloth_parse = self.parse_cloth(cloth_parse_path)
        # parsing
        if self.opt.warp_cloth:
            source_splitext = os.path.join(source_splitext.split('/')[0], source_splitext.split('/')[2])

        source_parse_vis_path = os.path.join('dataset/parse_cihp', source_splitext + '_vis.png')
        source_parse_vis = self.transforms['3'](Image.open(source_parse_vis_path))
        source_parse_path = os.path.join('dataset/parse_cihp', source_splitext + '.png')
        source_parse = pose_utils.parsing_embedding(source_parse_path)

        source_parse_shape = np.array(Image.open(source_parse_path))
        source_parse_shape = (source_parse_shape > 0).astype(np.float32)
        source_parse_shape = Image.fromarray((source_parse_shape*255).astype(np.uint8))
        source_parse_shape = source_parse_shape.resize((self.size[1]//16, self.size[0]//16), Image.BILINEAR) # downsample and then upsample
        source_parse_shape = source_parse_shape.resize((self.size[1], self.size[0]), Image.BILINEAR)
        source_parse_shape = self.transforms['1'](source_parse_shape) # [-1,1]

        source_parse_head = (np.array(Image.open(source_parse_path)) == 1).astype(np.float32) + \
                    (np.array(Image.open(source_parse_path)) == 2).astype(np.float32) + \
                    (np.array(Image.open(source_parse_path)) == 4).astype(np.float32) + \
                    (np.array(Image.open(source_parse_path)) == 13).astype(np.float32)

        # prepare for warped cloth
        phead = torch.from_numpy(source_parse_head) # [0,1]
        im_h = source_img * phead - (1 - phead) # [-1,1], fill -1 for other parts, thus become black visual

        # pose heatmap embedding
        source_pose_path = os.path.join('dataset/pose_coco', source_splitext +'_keypoints.json')
        with open(source_pose_path, 'r') as f:
            a = json.load(f)
            source_pose = a['people'][0]['pose_keypoints_2d']
        source_pose_loc = pose_utils.pose2loc(source_pose)
        source_pose_embedding = pose_utils.heatmap_embedding(self.size, source_pose_loc)

        if self.opt.warp_cloth:
            target_splitext = os.path.splitext(target_pose)[0]
            target_pose_path = os.path.join('dataset/pose_coco', data_suffix, target_splitext.split('/')[1] +'_keypoints.json')
            warped_cloth_name = source_splitext.split('/')[0] + '/' + \
                    source_splitext.split('/')[1] + '_' + \
                    target_splitext.split('/')[1] + '_' + \
                    cloth_splitext.split('/')[2] + '_warped_cloth.jpg'
            warped_cloth_name = warped_cloth_name.split('/')[0] + '/' + warped_cloth_name.split('/')[1].split('-')[0] + '/' + warped_cloth_name.split('/')[1]
        else:
            target_pose_path = os.path.join('dataset/pose_coco', target_pose)
            warped_cloth_name = cloth_splitext

        with open(target_pose_path, 'r') as f:
            a = json.load(f)
            target_pose = a['people'][0]['pose_keypoints_2d']
        target_pose_loc = pose_utils.pose2loc(target_pose)
        target_pose_embedding = pose_utils.heatmap_embedding(self.size, target_pose_loc)
        target_pose_img, _ = pose_utils.draw_pose_from_cords(target_pose_loc, (256, 192))
        result = {
                'source_parse': source_parse,
                'source_parse_vis': source_parse_vis,
                'source_pose_embedding': source_pose_embedding,
                'target_pose_embedding': target_pose_embedding,
                'target_pose_loc': target_pose_loc,
                'source_image': source_img,
                'cloth_image' : cloth_img,
                'cloth_parse' : cloth_parse,
                'source_parse_shape': source_parse_shape,
                'im_h': im_h, # source image head and hair
                'source_image_name': source_splitext,
                'cloth_image_name': cloth_splitext,
                'source_img_path': source_img_path,
                'target_pose_path': target_pose_path,
                'source_parse_vis_path': source_parse_vis_path,
                'target_pose_img': target_pose_img,
                'warped_cloth_name': warped_cloth_name
        }

        return result

    def __len__(self):

        return len(self.img_list)

    def make_pair(self, pair_list, test_num, pair_num):

        test_list = list(filter(lambda p: p.split('\t')[3] == 'test', pair_list))
        img_source = [i.split('\t')[0] for i in test_list]
        img_target = [i.split('\t')[1] for i in test_list]
        cloth_img = [i.split('\t')[2] for i in test_list]

        selected_img = random.sample(img_source, test_num)
        selected_target = random.sample(img_target, pair_num)
        selected_cloth = random.sample(cloth_img, pair_num)

        pair_list = []

        with open('demo/uniform_test.txt', 'w') as f:
            for i in range(test_num):
                for j in range(pair_num):
                    pair = selected_img[i] + '\t' + selected_target[j] + '\t' + selected_cloth[j] + '\t' + 'test'
                    f.write(pair + '\n')
                    pair_list.append(pair)
        return pair_list

    def open_transform(self, path, downsample=False):
        img = Image.open(path)
        if downsample:
            img = img.resize((96,128),Image.BICUBIC)
        img = self.transforms['3'](img)
        return img

    def parse_cloth(self, path, downsample=False):
        cloth_parse = Image.open(path)
        cloth_parse_array = np.array(cloth_parse)
        cloth_parse = (cloth_parse_array == 255).astype(np.float32) # 0 | 1
        cloth_parse = cloth_parse[np.newaxis, :]

        if downsample:
            [X, Y] = np.meshgrid(range(0,192,2),range(0,256,2))
            cloth_parse = cloth_parse[:,Y,X]

        cloth_parse = torch.from_numpy(cloth_parse)

        return cloth_parse

if __name__ == '__main__':
    pass


# -*- coding: utf-8 -*-
# coding=utf-8
# @author: cer
import numpy as np
import pickle

import data.load
import io
from metrics.accuracy import conlleval
from data_loader import *
import my_metrics as metric
from keras.models import Sequential,Model
from keras.layers.embeddings import Embedding
from keras.layers.recurrent import SimpleRNN, GRU, LSTM
from keras.layers.core import Dense, Dropout
from keras.layers.wrappers import TimeDistributed
from keras.layers import Convolution1D, MaxPooling1D, Bidirectional,Input, Embedding, LSTM, Dense
import progressbar
from keras.utils import plot_model
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score
import matplotlib.pyplot as plt
import requests
import speech_recognition as sr

with open('word2index.pkl', 'rb') as f:
    word2index = pickle.load(f)
with open('index2slot.pkl', 'rb') as f:
    index2slot =  pickle.load(f)
with open('index2intent.pkl', 'rb') as f:
    index2intent =  pickle.load(f)
with open('index2word.pkl', 'rb') as f:
    index2word = pickle.load(f)

### Model
n_classes = len(index2slot)
n_intent = len(index2intent)
n_vocab = len(index2word)
print(index2slot)
# print(index2word)

class NLU_Model():
    model = None
    def __init__(self):
        main_input = Input(name='main_input',shape=(None,))
        x = Embedding(output_dim=50, input_dim=n_vocab)(main_input)
        lstm_out = Bidirectional(LSTM(units=50,return_sequences=True))(x)
        lstm_out_slot = Bidirectional(LSTM(units=50,return_sequences=True))(lstm_out)
        intent_out = Bidirectional(LSTM(units=50,return_sequences=False))(lstm_out)
        main_out = TimeDistributed(Dense(n_classes, activation='softmax'))(lstm_out_slot)
        intent_out = Dense(n_intent, activation='softmax')(intent_out)
        model = Model(inputs= main_input , outputs= [main_out,intent_out])
        model.compile(optimizer='rmsprop', loss='categorical_crossentropy')
        print(model.summary())
        model.load_weights('best_model_weights.h5',by_name=True)
        self.model = model

    def NLU(self,input_utterance):
        input_token = input_utterance.split(" ")
        print(input_token)
        input_index = list(map(lambda i: word2index[i] if i in word2index else word2index["<UNK>"], input_token))
        print(input_index)
        input_index = np.array(input_index)
        input_index = input_index[np.newaxis, :]
        pred = self.model.predict_on_batch(input_index)
        slot_pred = np.argmax(np.array(pred[0]), -1)[0]
        slot_pred = list(map(lambda x: index2slot[x], slot_pred))
        print(slot_pred)
        return (slot_pred)

class RuleEngine:
    def execute(self,slots):
        if 'device_lamp' in slots:
            if "command_off" in slots:
                requests.get('http://192.168.1.242:82/led00')
            if "command_on" in slots:
                requests.get('http://192.168.1.242:82/led10')
            if "command_decrease" in slots:
                requests.get('http://192.168.1.242:82/led10')
            if "command_increase" in slots:
                requests.get('http://192.168.1.242:82/led11')
        if "device_tv" in slots:
            if "command_off" in slots:
                requests.get('http://192.168.1.242:82/ir_send1')
            if "command_on" in slots:
                requests.get('http://192.168.1.242:82/ir_send1')
        if "volume" in slots:
            if "command_decrease" in slots:
                requests.get('http://192.168.1.242:82/ir_send13')
                requests.get('http://192.168.1.242:82/ir_send13')
                requests.get('http://192.168.1.242:82/ir_send13')
            if "command_increase" in slots:
                requests.get('http://192.168.1.242:82/ir_send12')
                requests.get('http://192.168.1.242:82/ir_send12')
                requests.get('http://192.168.1.242:82/ir_send12')
        if "command_change_channel" in slots:
            if "command_decrease" in slots:
                requests.get('http://192.168.1.242:82/ir_send15')
            if "command_increase" in slots:
                requests.get('http://192.168.1.242:82/ir_send14')
            if "channel_num_1" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
            if "channel_num_2" in slots:
                requests.get('http://192.168.1.242:82/ir_send3')
            if "channel_num_3" in slots:
                requests.get('http://192.168.1.242:82/ir_send4')
            if "channel_num_4" in slots:
                requests.get('http://192.168.1.242:82/ir_send5')
            if "channel_num_5" in slots:
                requests.get('http://192.168.1.242:82/ir_send6')
            if "channel_num_6" in slots:
                requests.get('http://192.168.1.242:82/ir_send7')
            if "channel_num_7" in slots:
                requests.get('http://192.168.1.242:82/ir_send8')
            if "channel_num_8" in slots:
                requests.get('http://192.168.1.242:82/ir_send9')
            if "channel_num_9" in slots:
                requests.get('http://192.168.1.242:82/ir_send10')
            if "channel_num_10" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
                requests.get('http://192.168.1.242:82/ir_send11')
            if "channel_num_11" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
                requests.get('http://192.168.1.242:82/ir_send2')
            if "channel_num_12" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
                requests.get('http://192.168.1.242:82/ir_send3')
            if "channel_num_13" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
                requests.get('http://192.168.1.242:82/ir_send4')
            if "channel_num_14" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
                requests.get('http://192.168.1.242:82/ir_send5')
            if "channel_num_15" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
                requests.get('http://192.168.1.242:82/ir_send6')
            if "channel_num_16" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
                requests.get('http://192.168.1.242:82/ir_send7')
            if "channel_num_17" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
                requests.get('http://192.168.1.242:82/ir_send8')
            if "channel_num_18" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
                requests.get('http://192.168.1.242:82/ir_send9')
            if "channel_num_19" in slots:
                requests.get('http://192.168.1.242:82/ir_send2')
                requests.get('http://192.168.1.242:82/ir_send10')
            if "channel_num_20" in slots:
                requests.get('http://192.168.1.242:82/ir_send3')
                requests.get('http://192.168.1.242:82/ir_send11')


import pygame
import wx
from pygame.math import Vector2
from sw_serial import *
import cv2
import numpy as np
from datetime import date

# Initialize the serial connection and then joysticks
pygame.joystick.init()
joy_count = pygame.joystick.get_count()
if joy_count > 0:
    xbox_360 = pygame.joystick.Joystick(0)
    xbox_360.init()
    print(xbox_360.get_name())

###############################################################################
# Joystick controls
###############################################################################

def button_press():
    btn_dict = {
        0:BTN_A,1:BTN_B,2:BTN_X,3:BTN_Y,4:BTN_L,
        5:BTN_R,6:BTN_MINUS,7:BTN_PLUS,8:BTN_HOME}
    # Mapped from XBOX 360 to Switch
    btn_list = []
    for button in btn_dict:
        if xbox_360.get_button(button):
            btn_list.append(btn_dict[button])
    return (sum(btn_list))

def hat_motion():
    hat_xy = [[-1,0],[1,0],[0,1],[0,-1],[0,0]]
    # X, Y dpad values from XBOX controller
    hat_btn = [DPAD_L,DPAD_R,DPAD_U,DPAD_D,DPAD_CENTER]
    # Corresponding action in Switch controller
    for hat in hat_xy:
        if xbox_360.get_hat(0) == (hat[0],hat[1]):
            trigger_hat = hat_btn[hat_xy.index(hat)]
    return trigger_hat

def z_axis():
    if xbox_360.get_axis(2) >= .5:
        z_btn = BTN_ZL
    elif xbox_360.get_axis(2) <= -.5:
        z_btn = BTN_ZR
    else:
        z_btn = 0
    return z_btn

def stick_move():
    left_x = xbox_360.get_axis(0)
    left_y = xbox_360.get_axis(1)
    right_x = xbox_360.get_axis(4)
    right_y = xbox_360.get_axis(3)

    axis_list = [[left_x,left_y],[right_x,right_y]]

    angle_list = []
    for stick in axis_list:
        vec = Vector2(stick[0],stick[1])
        rad, ang = vec.as_polar()
        if ang <= 0:
            ang = (-1) * ang
        elif ang >= 0:
            ang = 360 - ang
        angle_list.append(ang)

 # Need to input intensity -> value between 50 (0x32) and 255 (0xFF).
 # rad adjusted to 1 if over 255 as it resets above this value.
    stick_intsy = []
    for stick in axis_list:
        power = Vector2(stick[0],stick[1])
        rad, ang = power.as_polar()
        if rad > 1:
            rad = 1
        intsy = int(rad * 255)
        stick_intsy.append(intsy)

# read the x and y from each stick and if > 0 consider pressed
# then send the stick signal
#if blank send center command (shared for L and R as it's the same value)
    state_axes = []
    for joy in axis_list:
        which_stick = axis_list.index(joy)
        state_joy = 0
        for axis in joy:
            if -1 <= axis <= -0.1 or 0.1 <= axis <= 1:
                state_joy += 1
            else:
                state_joy += 0
        if state_joy > 0:
            state_axes.append(
                stick_angle(int(angle_list[which_stick]),
                stick_intsy[which_stick], which_stick)
                )
        else:
            state_axes.append(LSTICK_CENTER)
    joy_cmd = state_axes[0] + state_axes[1]
    return joy_cmd

def timemov(direction, length, stick=0):
    i = 0
    while i < length:
        cmd = stick_angle(direction, 0xFF,stick)
        send_cmd(cmd)
        i += 1

def btn_press(button, wait):
    send_cmd(button)
    p_wait(0.1)
    send_cmd(BTN_NONE)
    p_wait(wait)

def stick_press(dpad,wait):
    send_cmd(dpad)
    p_wait(0.05)
    send_cmd(LSTICK_CENTER)
    p_wait(wait)

def hat_press(hat,wait):
    send_cmd(BTN_NONE + hat+ LSTICK_CENTER + RSTICK_CENTER)
    p_wait(0.07)
    send_cmd(DPAD_CENTER)
    p_wait(wait)

###############################################################################
# Keyboard inputs
###############################################################################
def kbd_input(key_name):
    if key_name == "0":
        print("Starting egg hatching macro...")
        hatch_egg_macro()
    if key_name == "2":
        print("Starting release macro...")
        release_macro()
    if key_name == "3":
        print("Next loop")
        encounter_loop()
    if key_name == "5":
        print("Starting batch hatching macro")
        batch_hatch_macro()
    if key_name == "7":
        print("Starting egg fetching macro...")
        get_eggs_macro()
    elif key_name == "w":
        cmd = stick_angle(90, 0xFF,0)
        send_cmd(cmd)
    elif key_name == "s":
        cmd = stick_angle(270, 0xFF,0)
        send_cmd(cmd)
    elif key_name == "a":
        cmd = stick_angle(180, 0xFF,0)
        send_cmd(cmd)
    elif key_name == "d":
        cmd = stick_angle(360, 0xFF,0)
        send_cmd(cmd)
    elif key_name == "q":
        send_cmd(RSTICK_L)
    elif key_name == "e":
        send_cmd(RSTICK_R)
    elif key_name == "up":
        hat_press(DPAD_U,0.01)
    elif key_name == "down":
        hat_press(DPAD_D,0.01)
    elif key_name == "left":
        hat_press(DPAD_L,0.01)
    elif key_name == "right":
        hat_press(DPAD_R,0.01)
    elif key_name == "return":
        btn_press(BTN_A,0.05)
    elif key_name == "backspace":
        btn_press(BTN_B,0.05)
    elif key_name == "x":
        btn_press(BTN_X,0.05)
    elif key_name == "y":
        btn_press(BTN_Y,0.05)
    elif key_name == "r":
        btn_press(BTN_R,0.05)
    elif key_name == "l":
        btn_press(BTN_L,0.05)
    elif key_name == "-":
        btn_press(BTN_MINUS,0.05)
    elif key_name == ";":
        btn_press(BTN_PLUS,0.05)
    elif key_name == "h":
        btn_press(BTN_HOME,0.05)
###############################################################################
# Macros
###############################################################################

class swmacro:
    app = wx.App() #One App per process, may need to move further up in future

    def __init__(self,msg,title,lower,upper):
        self.msg = msg
        self.title = title
        self.lower = lower
        self.upper = upper

    def macro_prompt(self):
        frame = wx.Frame(None, -1, 'win.py')
        frame.SetSize(0,0,100,100)
        dlg = wx.TextEntryDialog(frame, self.msg,self.title)
        if dlg.ShowModal() == wx.ID_OK:
            dlg.GetValue()
            int_input = int(dlg.GetValue())
            if int_input == None:
                return None
            if int_input in range(self.lower,self.upper+1):
                return int_input
            else:
                dialog = wx.MessageDialog(
                    frame,"Incorrect amount", 'Error',wx.OK|wx.ICON_ERROR
                    )
                dialog.ShowModal()
                dialog.Destroy()
                return swmacro.macro_prompt(self)

def hatch_egg_macro():
#Needs full party, town map first option, start w/o bike, egg ready to pick up
    egg_group = swmacro(
        'Choose Egg Group # (Starters(1) Longer(2) Very long(3): ',
        'Egg cycle type',1,3
        )
    group = egg_group.macro_prompt()
    egg_hatch = swmacro('Number to hatch','Egg hatch count',1,600)
    eggs = egg_hatch.macro_prompt()

    if egg_group == None or egg_hatch == None:
        print("Macro aborted due to no entry")
        return None

    steps_dict = {1:12800,2:14800,3:21200}
    try:
        steps = steps_dict.get(group)
        for i in range(eggs):
            btn_press(BTN_X,1) #Open menu - assuming map is first btn available
            btn_press(BTN_A,3) #Open map
            btn_press(BTN_A,1.5) #Select location assuming nursery is location
            btn_press(BTN_A,1) #Confirm
            btn_press(BTN_A,3) #Travel
            timemov(260,200) #Move back to egg person
            btn_press(BTN_L,1.3) #Swivel screen around to face
            btn_press(BTN_A,1.1) #open conversation
            btn_press(BTN_A,2.2) #next line
            btn_press(BTN_A,3.2) #"Yes" to accept egg
            btn_press(BTN_A,2.2) #Confirmation
            btn_press(BTN_A,1.3) #Add to roster
            btn_press(BTN_A,2.2) #Confirmation
            stick_press(LSTICK_D,0.7) #Down 1 slot
            stick_press(LSTICK_D,0.7) #Down another slot (now on 3rd slot)
            btn_press(BTN_A,2.5) #Confirmation
            btn_press(BTN_A,2.5) #Confirmation
            btn_press(BTN_A,2) #Confirmation
            send_cmd(stick_angle(270, 0xFF,0))#turn around
            p_wait(0.3)
        #        btn_press(BTN_PLUS,1) #Get on bike
            timemov(250,700) #Move into place
            btn_press(BTN_PLUS,1.2)
            # 360 Circle @ Full Intensity
            for r in range(steps):
                send_cmd(LSTICK_L+RSTICK_R)
            btn_press(BTN_A,18)#After egg hatches
            btn_press(BTN_A,5)#After egg hatches
            btn_press(BTN_PLUS,2.3)
            print(i)
    except:
        print("Incorrect Group")

def get_eggs_macro():
    get_egg = swmacro('Number of eggs to get','Egg fetching',1,150)
    egg_count = get_egg.macro_prompt()
    if egg_count == None:
        print("Macro aborted due to no entry")
        return None
    print("Getting %s eggs" % egg_count)
    #Needs to have egg ready. full slots
    for g in range (egg_count):
            btn_press(BTN_X,1) #Open menu - assuming map is first btn available
            btn_press(BTN_A,3) #Open map
            btn_press(BTN_A,1.5) #Select location assuming nursery is location
            btn_press(BTN_A,1) #Confirm
            btn_press(BTN_A,3) #Travel
            timemov(260,200) #Move back to egg person
            btn_press(BTN_L,1) #Swivel screen around to face
            btn_press(BTN_A,1) #open conversation
            btn_press(BTN_A,2) #next line
            btn_press(BTN_A,3.2) #"Yes" to accept egg
            btn_press(BTN_A,2) #Confirmation
            stick_press(LSTICK_D,0.5) #Down 1 slot
            btn_press(BTN_A,1.5) #Confirm
            btn_press(BTN_A,2) #Confirm
            btn_press(BTN_A,1) #exit conversation
            send_cmd(stick_angle(270, 0xFF,0))#turn around
            p_wait(0.2)
            timemov(250,700) #Move into place
            btn_press(BTN_PLUS,1)
            # 360 Circle @ Full Intensity
            r = 0
            for r in range (3000):
                #estimated hatch rate 256 steps x 5 sets at 80% rate
                # and using ratio of steps with egg hatching
                send_cmd(LSTICK_L+RSTICK_R)
            btn_press(BTN_PLUS,1)
   #return egg_count

def batch_hatch_macro():
    steps = 13000
    #Start at current inventory (2nd position down)
    #Eggs need to be in second column, with blank column at first
    btn_press(BTN_Y,1) #Change cursor
    btn_press(BTN_Y,1) #Change again
    btn_press(BTN_A,1) #Select first item
    hat_press(DPAD_D,0.3)
    hat_press(DPAD_D,0.3)
    hat_press(DPAD_D,0.3)
    hat_press(DPAD_D,0.3)#Move down to fifth item
    btn_press(BTN_A,1)#Select and pick up
    hat_press(DPAD_R,0.5)#Move over to inventory
    hat_press(DPAD_U,0.5) #Shift up
    btn_press(BTN_A,1) #Set down in box

    hat_press(DPAD_R,0.3) #Move to row of eggs
    btn_press(BTN_A,1) #Select first egg
    hat_press(DPAD_D,0.3)
    hat_press(DPAD_D,0.3)
    hat_press(DPAD_D,0.3)
    hat_press(DPAD_D,0.3)#Move down to fifth egg
    btn_press(BTN_A,1) #Select and pick up
    hat_press(DPAD_L,0.5)
    hat_press(DPAD_L,0.5)#Move to the てもち list
    hat_press(DPAD_D,0.3)#Position for placement
    btn_press(BTN_A,1) #Place in active roster
    btn_press(BTN_B,1.5)
    btn_press(BTN_B,1.5)
    btn_press(BTN_B,1.5) #get back to overworld
    btn_press(BTN_PLUS,1) #Get on bike
    for r in range(steps):
        send_cmd(LSTICK_L+RSTICK_R)
    for i in range(4):
        btn_press(BTN_A,18)#After egg hatches
        btn_press(BTN_A,6)#After egg hatches
        #Timing is currently INCORRECT, needs to be a bit longer and test
        if i != 4:
            for r in range(200):
                send_cmd(LSTICK_L+RSTICK_R)
        else:
            btn_press(BTN_PLUS,2.3)

def encounter_loop():

    i = 0
    try:
        while True:
        #Start in battle in order to shiny check
            btn_press(BTN_HOME,0.7) #Go to home screen
            btn_press(BTN_X,0.8) #bring up exit menu
            btn_press(BTN_A,2.8) #select exit
            btn_press(BTN_A,1) #bring up profile select
            btn_press(BTN_A,17.5) #open game
            btn_press(BTN_A,6.5) #load game
            btn_press(BTN_A,5) #open menu
            btn_press(BTN_A,1.5) #confirm
            #btn_press(BTN_A,1.1) #confirm (Regidrago)
            #btn_press(BTN_A,1.5) #confirm (Regidrago)
            btn_press(BTN_A,13.5) #start encounter
            #set cv2 for screenshot
            cap = cv2.VideoCapture(0,cv2.CAP_DSHOW)
            cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
            ret, frame = cap.read()
            img_hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            mask_color1 = np.asarray([14,0,50])
            mask_color2 = np.asarray([24,255,255])
            mask = cv2.inRange(img_hsv, mask_color1,mask_color2)
            color_pct = (cv2.countNonZero(mask)/(img_hsv.shape[0] * img_hsv.shape[1]))*100
            print(str(i)+' {:.5f}'.format(color_pct))
            if i == 0:
                if color_pct < 1.5:
                    cv2.imshow("Shiny",shiny_pic)
                    break
            if i != 0:
                if color_pct < 1.6:
                    cv2.imshow("Shiny",shiny_pic)
                    break
            p_wait(3)
            i += 1
    except KeyboardInterrupt:
        pass

def release_macro():
    #Need to be at start position of box for deleting correctly
    #Look into adding support for an arbitrary start point
    release_pokemon = swmacro(
        'Number of Pokemon to release',
        'Pokemon Count',1,600
        )
    release_count = release_pokemon.macro_prompt()
    if release_count == None:
        print("Macro aborted due to no entry")
        return None
    r = 1
    b = 1
    try:
        for i in range(release_count):
            btn_press(BTN_A,0.5)
            hat_press(DPAD_U,0.2)
            hat_press(DPAD_U,0.2)
            btn_press(BTN_A,0.7)
            hat_press(DPAD_U,0.2)
            btn_press(BTN_A,1.2)
            btn_press(BTN_A,0.9)
            print("Released %s" % i)
            if  (i+1) % 6 == 0 and (i+1) % 30 != 0:
                #Added 1 to account for r starting at 0
                hat_press(DPAD_D,0.8)
                print("%s end of row" % (i+1))
                r += 1
                print(r)
            elif (i+1) % 30 == 0:
                hat_press(DPAD_D,0.3)
                hat_press(DPAD_D,0.3)
                hat_press(DPAD_D,0.3)
                hat_press(DPAD_R,0.3)
                hat_press(DPAD_R,0.3)
                btn_press(BTN_R,0.8)
                print("%s end of box" % (i+1))
                b += 1
                r += 1
                print(b)
            else:
                    if (b % 2 != 0 and r % 2 == 0) or (r % 2 != 0 and b % 2 == 0):
                        hat_press(DPAD_L,0.3)
                        print("%s next value (left)" % (i+1))
                    else:
                        hat_press(DPAD_R,0.3)
                        print("%s next value (right)" % (i+1))
    except:
        print("Incorrect value")

def watt_farmer():
    i = 0
    day = 1
    while i < 1000:
        print(f"Current day is {day}")
        if day == 1:
            btn_press(BTN_A,1.1)
            btn_press(BTN_A,3)
        else:
            btn_press(BTN_A,1)
            btn_press(BTN_A,1)
            btn_press(BTN_A,1.5)
            btn_press(BTN_A,3)
        # Enter battle search for first time (must be facing den)
        btn_press(BTN_HOME,1)
        hat_press(DPAD_R,0.3)
        hat_press(DPAD_R,0.3)
        hat_press(DPAD_D,0.4)
        hat_press(DPAD_R,0.3)
        btn_press(BTN_A,1)
        # Brings to settings screen
        for r in range (14):
            hat_press(DPAD_D,0.2)
        # To quickly scroll down the list
        hat_press(DPAD_R,0.3)
        hat_press(DPAD_D,0.3)
        hat_press(DPAD_D,0.3)
        hat_press(DPAD_D,0.3)
        hat_press(DPAD_D,0.3)
        btn_press(BTN_A,0.6)
        hat_press(DPAD_D,0.3)
        hat_press(DPAD_D,0.3)
        btn_press(BTN_A,0.6)
        # Brings up time settings
        hat_press(DPAD_R,0.2)
        hat_press(DPAD_R,0.2)
        hat_press(DPAD_U,0.2)
        btn_press(BTN_A,0.2)
        btn_press(BTN_A,0.2)
        btn_press(BTN_A,0.2)
        btn_press(BTN_A,0.2)
        btn_press(BTN_HOME,1.4)
        # Moves date up and moves back to home screen
        btn_press(BTN_A,1.5)
        btn_press(BTN_B,1)
        btn_press(BTN_A,4)
        # Exit battle
        if day == 31:
            day = 1
        else:
            day += 1
        i += 1


"""Extracting Content-Aware Perceptual Features using Pre-Trained ResNet-50"""
# Author: Dingquan Li
# Email: dingquanli AT pku DOT edu DOT cn
# Date: 2018/3/27
#
# CUDA_VISIBLE_DEVICES=0 python CNNfeatures.py --database=KoNViD-1k --frame_batch_size=64
# CUDA_VISIBLE_DEVICES=1 python CNNfeatures.py --database=CVD2014 --frame_batch_size=32
# CUDA_VISIBLE_DEVICES=0 python CNNfeatures.py --database=LIVE-Qualcomm --frame_batch_size=8

import torch
from torchvision import transforms, models
import torch.nn as nn
from torch.utils.data import Dataset
import skvideo.io
from PIL import Image
import os
#import h5py
import numpy as np
import random
from argparse import ArgumentParser

class VideoDataset(Dataset):
    """Read data from the original dataset for feature extraction"""
    def __init__(self, videos_dir, video_names, score, video_format='RGB', width=None, height=None):

        super(VideoDataset, self).__init__()
        self.videos_dir = videos_dir
        self.video_names = video_names
        self.score = score
        self.format = video_format
        self.width = width
        self.height = height

    def __len__(self):
        return len(self.video_names)

    def __getitem__(self, idx):
        video_name = self.video_names[idx]
        assert self.format == 'YUV420' or self.format == 'RGB'
        if self.format == 'YUV420':
            video_data = skvideo.io.vread(os.path.join(self.videos_dir, video_name), self.height, self.width, inputdict={'-pix_fmt':'yuvj420p'})
        else:
            video_data = skvideo.io.vread(os.path.join(self.videos_dir, video_name))
        video_score = self.score[idx]

        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

        video_length = video_data.shape[0]
        video_channel = video_data.shape[3]
        video_height = video_data.shape[1]
        video_width = video_data.shape[2]
        transformed_video = torch.zeros([video_length, video_channel,  video_height, video_width])
        for frame_idx in range(video_length):
            frame = video_data[frame_idx]
            frame = Image.fromarray(frame)
            frame = transform(frame)
            transformed_video[frame_idx] = frame

        sample = {'video': transformed_video,
                  'score': video_score}

        return sample

class ResNet50(torch.nn.Module):
    """Modified ResNet50 for feature extraction"""
    def __init__(self):
        super(ResNet50, self).__init__()
        self.features = nn.Sequential(*list(models.resnet50(pretrained=True).children())[:-2])
        for p in self.features.parameters():
            p.requires_grad = False

    def forward(self, x):
        # features@: 7->res5c
        for ii, model in enumerate(self.features):
            x = model(x)
            if ii == 7:
                features_mean = nn.functional.adaptive_avg_pool2d(x, 1)
                features_std = global_std_pool2d(x)
                return features_mean, features_std

def global_std_pool2d(x):
    """2D global standard variation pooling"""
    return torch.std(x.view(x.size()[0], x.size()[1], -1, 1),
                     dim=2, keepdim=True)

def get_features(video_data, frame_batch_size=64, device='cuda'):
    """feature extraction"""
    extractor = ResNet50().to(device)
    video_length = video_data.shape[0]
    frame_start = 0
    frame_end = frame_start + frame_batch_size
    output1 = torch.Tensor().to(device)
    output2 = torch.Tensor().to(device)
    extractor.eval()
    with torch.no_grad():
	    while frame_end < video_length:
	        batch = video_data[frame_start:frame_end].to(device)
	        features_mean, features_std = extractor(batch)
	        output1 = torch.cat((output1, features_mean), 0)
	        output2 = torch.cat((output2, features_std), 0)
	        frame_end += frame_batch_size
	        frame_start += frame_batch_size

	    last_batch = video_data[frame_start:video_length].to(device)
	    features_mean, features_std = extractor(last_batch)
	    output1 = torch.cat((output1, features_mean), 0)
	    output2 = torch.cat((output2, features_std), 0)
	    output = torch.cat((output1, output2), 1).squeeze()

    return output

if __name__ == "__main__":
    parser = ArgumentParser(description='"Extracting Content-Aware Perceptual Features using Pre-Trained ResNet-50')
    parser.add_argument("--seed", type=int, default=19920517)
    parser.add_argument('--database', default='KoNViD-1k', type=str,
                        help='database name (default: KoNViD-1k)')
    parser.add_argument('--frame_batch_size', type=int, default=64,
                        help='frame batch size for feature extraction (default: 64)')

    parser.add_argument('--disable_gpu', action='store_true',
                        help='flag whether to disable GPU')
    args = parser.parse_args()

    torch.manual_seed(args.seed)  #
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(args.seed)
    random.seed(args.seed)

    torch.utils.backcompat.broadcast_warning.enabled = True

    if args.database == 'KoNViD-1k':
        videos_dir = '/home/ldq/Downloads/KoNViD-1k/'  # videos dir
        features_dir = 'CNN_features_KoNViD-1k/'  # features dir
        datainfo = 'data/KoNViD-1kinfo.mat'  # database info: video_names, scores; video format, width, height, index, ref_ids, max_len, etc.
    if args.database == 'CVD2014':
        videos_dir = '/media/ldq/Research/Data/CVD2014/'
        features_dir = 'CNN_features_CVD2014/'
        datainfo = 'data/CVD2014info.mat'
    if args.database == 'LIVE-Qualcomm':
        videos_dir = '/media/ldq/Others/Data/12.LIVE-Qualcomm Mobile In-Capture Video Quality Database/'
        features_dir = 'CNN_features_LIVE-Qualcomm/'
        datainfo = 'data/LIVE-Qualcomminfo.mat'

    if not os.path.exists(features_dir):
        os.makedirs(features_dir)

    device = torch.device("cuda" if not args.disable_gpu and torch.cuda.is_available() else "cpu")

    Info = h5py.File(datainfo, 'r')
    video_names = [Info[Info['video_names'][0, :][i]][()].tobytes()[::2].decode() for i in range(len(Info['video_names'][0, :]))]
    scores = Info['scores'][0, :]
    video_format = Info['video_format'][()].tobytes()[::2].decode()
    width = int(Info['width'][0])
    height = int(Info['height'][0])
    dataset = VideoDataset(videos_dir, video_names, scores, video_format, width, height)

    for i in range(len(dataset)):
        current_data = dataset[i]
        current_video = current_data['video']
        current_score = current_data['score']
        print('Video {}: length {}'.format(i, current_video.shape[0]))
        features = get_features(current_video, args.frame_batch_size, device)
        np.save(features_dir + str(i) + '_resnet-50_res5c', features.to('cpu').numpy())
        np.save(features_dir + str(i) + '_score', current_score)