from igraph import *
import numpy as np
import pickle
import json
from api_pkg.utils.request import *
from api_pkg.utils.metrics import *
from api_pkg.utils.tokenization import *
from api_pkg import dandelion, dbspotlight, opencalais, babelfy, adel, meaning_cloud, alchemy, textrazor
from itertools import combinations
from langdetect import detect
import time

EMBEDDING_DATA_PATH = 'data/embedding_data/'
WIKIDATA_CLASSES = pickle.load(open(EMBEDDING_DATA_PATH + 'all_classes.p', 'rb'))
WIKIDATA_DICT = json.load(open(EMBEDDING_DATA_PATH + 'wikidata_types_mapping.json'))
VIRTUAL_NODES = json.load(open(EMBEDDING_DATA_PATH + 'vitual_nodes.json'))
PATH_MAPPING_TYPES = "data/types_representation/types_mapping.json"

MAPPING_TYPES_OBJ = json.load(open(PATH_MAPPING_TYPES))

def getGraphNodes_and_Properties():
    embdedding_specifications_dict = json.load(open(EMBEDDING_DATA_PATH + 'embedding_specifications.json'))
    graph_path = embdedding_specifications_dict['structural']
    graph_obj = pickle.load(open(graph_path, 'rb'))
    MAX_DISTANCE_NODES = graph_obj['max_distance']
    structural_graph = graph_obj['graph']
    properties_list = embdedding_specifications_dict['semantic']
    return structural_graph, properties_list, MAX_DISTANCE_NODES

structural_graph, properties_list, MAX_DISTANCE_NODES = getGraphNodes_and_Properties()

def getWikidataTypeRepresentation(types_ordered_list):
    features_array = list()
    n_features = len(WIKIDATA_DICT.keys())
    resource_dict_representation = {t: [0 for i in range(n_features)] + [int(t not in WIKIDATA_CLASSES)] for t in
                                    (set(types_ordered_list) - {np.NAN})}
    for i, tag in enumerate(WIKIDATA_DICT.keys()):
        set_intersection = set(WIKIDATA_DICT[tag]) & set(resource_dict_representation.keys())
        for r in set_intersection:
            resource_dict_representation[r][i] = 1
    for t in types_ordered_list:
        if t in resource_dict_representation:
            features_array.append(resource_dict_representation[t])
        else:
            features_array.append([0 for i in range(n_features)] + [0])
    return np.array(features_array)

def getTypeRepresentation(types_ordered_list, ontology_name, mapping_types_obj=MAPPING_TYPES_OBJ, raise_flag=True):
    if ontology_name == 'wikidata':
        return getWikidataTypeRepresentation(types_ordered_list)
    else:
        features_array = list()
        ontology_path = MAPPING_TYPES_OBJ[ontology_name]
        ontology_representation = json.load(open(ontology_path))
        n_features = len(list(ontology_representation.values())[0])
        for t_info in types_ordered_list:
            vector = np.zeros(n_features)
            if type(t_info) == str:
                for t in t_info.split(','):
                    if t in ontology_representation:
                        vector_p = ontology_representation[t]
                        vector = [min([x + y, 1]) for x, y in zip(vector, vector_p)]
                    elif '-1' in ontology_representation:
                        print('AIA', ontology_name, t)
                        vector_p = ontology_representation['-1']
                        vector = [min([x + y, 1]) for x, y in zip(vector, vector_p)]
            features_array.append(vector)
        return np.array(features_array)

'''

    types_ordered_list1 = [a['uri'] for a in annotations]
    type_features1 = getTypeRepresentation(types_ordered_list1,self.ontology_uri)
    types_ordered_list2 = [a['type'] for a in annotations]
    type_features2 = getTypeRepresentation(types_ordered_list2,self.ontology_type)
    type_features = np.append(type_features1, type_features2, axis=1)

'''

def getTypeFeatures(type_list_dict, entity_list_dict, ontology_type_dict, ontology_entity_dict):
    type_features_obj = {}
    extractors_names = set(type_list_dict.keys()) | set(entity_list_dict.keys())
    for ext in extractors_names:
        if ext in type_list_dict and ext in entity_list_dict:
            type_features_obj[ext] = np.append(getTypeRepresentation(type_list_dict[ext], ontology_type_dict[ext]),
                                               getTypeRepresentation(entity_list_dict[ext], ontology_entity_dict[ext]),
                                               axis=1)
        elif ext in type_list_dict:
            type_features_obj[ext] = getTypeRepresentation(type_list_dict[ext], ontology_type_dict[ext])
        elif ext in entity_list_dict:
            type_features_obj[ext] = getTypeRepresentation(entity_list_dict[ext], ontology_entity_dict[ext])
    return type_features_obj

def getDistanceNodes(wd_id_1, wd_id_2, structural_graph):
    virtuals = 0
    try:
        n1 = VIRTUAL_NODES[n1]
        virtuals += 1
    except:
        n1 = wd_id_1
    try:
        n2 = VIRTUAL_NODES[n2]
        virtuals += 1
    except:
        n2 = wd_id_2
    if virtuals == 2:
        if wd_id_1 != wd_id_2:
            dist = 2
        else:
            dist = 0
    else:
        dist = structural_graph.shortest_paths(n1, n2)[0][0] + virtuals
    return dist

def getStructSimilarity(wd_id_1, wd_id_2, structural_graph, MAX_DISTANCE_NODES):
    if type(wd_id_1) != str or type(wd_id_2) != str:
        return 0.0
    else:
        try:
            d_struct = min([1.0, getDistanceNodes(wd_id_1, wd_id_2, structural_graph) / MAX_DISTANCE_NODES])
        except:
            return 0.0
        sim_struct = abs(1.0 - d_struct)
        return sim_struct

def assignMetric(s):
    if s == "one-hot":
        return oneHotSimilarity
    elif s == 'fuzzy':
        return fuzzSimilarity
    elif s == 'tf-idf':
        return ffIdfSimilarity

def getLabelsFeaturesToken(labels, metric=fuzzSimilarity):
    label_features = []
    for t in combinations(labels, 2):
        l1, l2 = t
        label_features.append(metric(l1, l2))
    return label_features

def getSematicSimilarity(wd_id_1, wd_id_2, properties, lang):
    semantic_similarity_array = []
    for p in properties:
        if 'filter' in p:
            items_1 = set()
            items_2 = set()
            if lang in p['filter']:
                items_1 = items_1 | set(wikiQuery(wd_id_1, p['property'], filter_q=p['filter'][lang])['o'])
                items_2 = items_2 | set(wikiQuery(wd_id_2, p['property'], filter_q=p['filter'][lang])['o'])
        else:
            items_1 = set(wikiQuery(wd_id_1, p['property'])['o'])
            items_2 = set(wikiQuery(wd_id_2, p['property'])['o'])
        sim_func = assignMetric(p['metric'])
        similarities = [sim_func(i1, i2) for i2 in items_2 for i1 in items_1]
        if len(similarities) != 0:
            semantic_similarity_array.append(max(similarities))
        else:
            semantic_similarity_array.append(0.0)
    return semantic_similarity_array

def getUrisSimilarityVector(wd_id_1, wd_id_2, G=structural_graph, properties_list=properties_list,
                            MAX_DISTANCE_NODES=MAX_DISTANCE_NODES, lang='fr'):
    if type(wd_id_1) == str and type(wd_id_2) == str:
        sim_struct = getStructSimilarity(wd_id_1, wd_id_2, structural_graph, MAX_DISTANCE_NODES)
        sim_semantic = getSematicSimilarity(wd_id_1, wd_id_2, properties_list, lang)
        return np.array([sim_struct] + sim_semantic)
    else:
        return np.zeros(len(properties_list) + 1)

def getEntityFeatures(entity_list_dict, ontology_entity_dict, lang):
    PRECOMPUTED_SIM = {}
    extractors_uris = list(entity_list_dict.keys())
    com_ext = [sorted(item) for item in list(combinations(extractors_uris, 2))] + [(ext, ext) for ext in
                                                                                extractors_uris]

    entity_features = {ext1: {ext2: [] for ext2 in extractors_uris} for ext1 in extractors_uris}
    entity_MATRIX = dict()

    length_list = len(entity_list_dict[extractors_uris[0]])
    for k in range(length_list):
        for comb in com_ext:
            uri1, uri2 = entity_list_dict[comb[0]][k], entity_list_dict[comb[1]][k]
            if not bool(uri1):
                uri1 = np.NAN
            if not bool(uri2):
                uri2 = np.NAN

            key = str(uri1) + '_' + str(uri2) + '_' + lang
            try:
                sim = PRECOMPUTED_SIM[key]
            except:
                sim = list(getUrisSimilarityVector(uri1, uri2, lang=lang)) + [
                    int(type(uri1) == type(uri2) == float or uri1 == uri2)]
                PRECOMPUTED_SIM[key] = sim
            if (uri1, uri2) not in entity_MATRIX:
                entity_MATRIX[(uri1, uri2)] = sim
            if (uri2, uri1) not in entity_MATRIX:
                entity_MATRIX[(uri2, uri1)] = sim
            entity_features[comb[0]][comb[1]].append(sim)
            if comb[0] != comb[1]:
                entity_features[comb[1]][comb[0]].append(sim)

    for ext in entity_features:
        for z, key in enumerate(extractors_uris):
            if z != 0:
                X_ext = np.append(X_ext, entity_features[ext][key], axis=1)
            else:
                X_ext = np.array(entity_features[ext][key])
        entity_features[ext] = X_ext
    return entity_features, entity_MATRIX

def fromOntologyEdglist_to_ClassRapresentation_notexclusive(edgelist_pd, set_root, limit):
    classes_rapresentation_dict = dict()
    roots = deepcopy(set_root)
    total_len_rapresentation = 0
    flag = False
    while len(classes_rapresentation_dict) != limit:
        df = edgelist_pd[edgelist_pd["class"].isin(roots)]
        roots = set(df['subclass'])
        len_rapresentation = len(roots)
        for i, r in enumerate(roots):
            past_features = []
            if flag:
                to_add_list = [classes_rapresentation_dict[c] for c in set(df[df['subclass'] == r]['class'])]
                for to_add in to_add_list:
                    if len(past_features) != 0:
                        past_features = [x + y for x, y in zip(past_features, to_add)]
                    else:
                        past_features = to_add
            classes_rapresentation_dict[r] = past_features + [int(i == j) for j in range(len_rapresentation)]
        flag = True

        for cl in (set(classes_rapresentation_dict.keys()) - roots):
            classes_rapresentation_dict[cl] += [0 for i in range(len_rapresentation)]
    return classes_rapresentation_dict

def fromOntologyEdglist_to_ClassRapresentation(edgelist_pd):
    subclasses = set(edgelist_pd["subclass"])
    classes = set(edgelist_pd["class"])
    set_root = classes - subclasses
    if len(set_root) != 1:
        for r in set_root:
            edgelist_pd = edgelist_pd.append([{
                "class": "_root_",
                "subclass": r
            }]
                , ignore_index=True)
        subclasses = set(edgelist_pd["subclass"])
        classes = set(edgelist_pd["class"])
        set_root = classes - subclasses
    return fromOntologyEdglist_to_ClassRapresentation_notexclusive(edgelist_pd, set_root, len(classes | subclasses) - 1)

from pyfasttext import FastText

FASTTEXT_EN = FastText('data/fasttext_data/wiki.en.bin')
FASTTEXT_FR1 = FastText('data/fasttext_data/wiki.fr.bin')
FASTTEXT_FR2 = FastText('data/fasttext_data/my_corpus_model.bin')

def getFastTextFeatures(text, lang):
    tokens = splitInTokens(text)
    if lang == 'fr':
        fasttextfeatures = np.array([
            np.append(FASTTEXT_FR1[token[0]], FASTTEXT_FR2[token[0]])
            for i, token in enumerate(tokens)])
    if lang == 'en':
        fasttextfeatures = np.array([FASTTEXT_EN[token[0]] for i, token in enumerate(tokens)])
    return fasttextfeatures

def getAnnotations(text, lang=None, model_setting='default'):
    extractors_list = [
        alchemy.ALCHEMY(),
        adel.ADEL(),
        dbspotlight.DBSPOTLIGHT(),
        opencalais.OPENCALAIS(),
        meaning_cloud.MEANINGCLOUD(),
        dandelion.DANDELION(),
        babelfy.BABELFY(),
        textrazor.TEXTRAZOR()
    ]

    # print(strftime("%H:%M:%S", gmtime()))
    limit_failures = 3
    waiting_secs = 7
    print('Request phase')
    for ext in extractors_list:
        print(ext.name)
        counter_failures = 0
        while counter_failures >= 0 and counter_failures < limit_failures:
            try:
                if ext.name == 'adel':
                    ext.extract(text, lang=lang, setting=model_setting)
                else:
                    ext.extract(text, lang=lang)
                counter_failures = -1
            except:
                counter_failures += 1
                if counter_failures == limit_failures:
                    raise Exception("The extractor", ext.name, "presented an error during the API request phase\n"
                                    + str(sys.exc_info()[1]))
                else:
                    time.sleep(waiting_secs)

    # print(strftime("%H:%M:%S", gmtime()))
    extractors_responses = {ext.name: ext.get_annotations() for ext in extractors_list}
    print('Parsing phase')
    for ext in extractors_list:
        print(ext.name)
        try:
            ext.parse()
        except:
            print(ext.get_annotations())
            raise Exception("The extractor", ext.name, "presented an error during the API response parsing phase\n" +
                            str(sys.exc_info()[1]))
    # print(strftime("%H:%M:%S", gmtime()))
    print('Tokenization phase')
    for ext in extractors_list:
        try:
            ext.tokenize()
        except:
            print("The extractor", ext.name, "presented an error during the API response tokenizing phase")
            print(sys.exc_info()[1])

    # print(strftime("%H:%M:%S", gmtime()))
    type_list_dict = {}
    entity_list_dict = {}
    score_list_dict = {}
    ontology_type_dict = {}

    ontology_entity_dict = {}
    for ext in extractors_list:
        annotations_ext = ext.get_annotations()
        try:
            ontology_type_dict[ext.name] = ext.ontology
            ontology_entity_dict[ext.name] = ext.ontology
        except:
            ontology_type_dict[ext.name] = ext.ontology_type
            ontology_entity_dict[ext.name] = ext.ontology_uri

        if ext.recognition:
            type_list_dict[ext.name] = [a['type'] for a in annotations_ext]
        if ext.disambiguation:
            entity_list_dict[ext.name] = [a['uri'] for a in annotations_ext]

        ext_scores = []
        for i, a in enumerate(annotations_ext):
            scores_list = list()
            try:
                scores_list.append(a['relevance'])
            except:
                pass
            try:
                scores_list.append(a['confidence'])
            except:
                pass
            if i == 0 and not bool(scores_list):
                break
            ext_scores.append(scores_list)

        if bool(ext_scores):
            score_list_dict[ext.name] = np.array(ext_scores)

    return extractors_responses, type_list_dict, score_list_dict, entity_list_dict, ontology_type_dict, ontology_entity_dict

'''

'''

def getFeatures(text, lang=None, model_setting='default'):
    if not bool(lang):
        lang = detect(text)
    features_dict_all = {'features': {}}
    extractors_responses, type_list_dict, score_list_dict, entity_list_dict, ontology_type_dict, ontology_entity_dict = getAnnotations(
        text, lang=lang, model_setting=model_setting)
    features_dict_all["entity_list"] = entity_list_dict
    features_dict_all["type_list"] = type_list_dict
    features_dict_all["extractors_responses"] = extractors_responses
    print('Forming features')

    if 'fasttext' not in features_dict_all['features']:
        features_dict_all['features']['fasttext'] = getFastTextFeatures(text, lang)
    print('Formed Fasttext features')

    if 'type' not in features_dict_all['features']:
        features_dict_all['features']['type'] = getTypeFeatures(type_list_dict, entity_list_dict, ontology_type_dict,
                                                                ontology_entity_dict)
    print('Formed type features')
    if 'score' not in features_dict_all['features']:
        features_dict_all['features']['score'] = score_list_dict

    print('Formed score features')
    if 'entity' not in features_dict_all['features']:
        features_dict_all['features']['entity'], features_dict_all['features']['entity_MATRIX'] = getEntityFeatures(
            entity_list_dict, ontology_entity_dict, lang)
    print('Formed entity features')
    return features_dict_all


from __future__ import absolute_import, division, print_function

import numpy as np
import torch
import torch.nn as nn

from .resnet import *
from .mobilenet import *
from .layers import Conv3x3, ConvBlock, upsample, Cube2Equirec, Concat, BiProj, CEELayer

from collections import OrderedDict

class UniFuse(nn.Module):
    """ UniFuse Model: Resnet based Euqi Encoder and Cube Encoder + Euqi Decoder
    """
    def __init__(self, num_layers, equi_h, equi_w, pretrained=False, max_depth=10.0,
                 fusion_type="cee", se_in_fusion=True):
        super(UniFuse, self).__init__()

        self.num_layers = num_layers
        self.equi_h = equi_h
        self.equi_w = equi_w
        self.cube_h = equi_h//2

        self.fusion_type = fusion_type
        self.se_in_fusion = se_in_fusion

        # encoder
        encoder = {2: mobilenet_v2,
                   18: resnet18,
                   34: resnet34,
                   50: resnet50,
                   101: resnet101,
                   152: resnet152}

        if num_layers not in encoder:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))
        self.equi_encoder = encoder[num_layers](pretrained)
        self.cube_encoder = encoder[num_layers](pretrained)

        self.num_ch_enc = np.array([64, 64, 128, 256, 512])
        if num_layers > 34:
            self.num_ch_enc[1:] *= 4

        if num_layers < 18:
            self.num_ch_enc = np.array([16, 24, 32, 96, 320])

        # decoder
        self.num_ch_dec = np.array([16, 32, 64, 128, 256])
        self.equi_dec_convs = OrderedDict()
        self.c2e = {}

        Fusion_dict = {"cat": Concat,
                       "biproj": BiProj,
                       "cee": CEELayer}
        FusionLayer = Fusion_dict[self.fusion_type]

        self.c2e["5"] = Cube2Equirec(self.cube_h // 32, self.equi_h // 32, self.equi_w // 32)

        self.equi_dec_convs["fusion_5"] = FusionLayer(self.num_ch_enc[4], SE=self.se_in_fusion)
        self.equi_dec_convs["upconv_5"] = ConvBlock(self.num_ch_enc[4], self.num_ch_dec[4])

        self.c2e["4"] = Cube2Equirec(self.cube_h // 16, self.equi_h // 16, self.equi_w // 16)
        self.equi_dec_convs["fusion_4"] = FusionLayer(self.num_ch_enc[3], SE=self.se_in_fusion)
        self.equi_dec_convs["deconv_4"] = ConvBlock(self.num_ch_dec[4] + self.num_ch_enc[3], self.num_ch_dec[4])
        self.equi_dec_convs["upconv_4"] = ConvBlock(self.num_ch_dec[4], self.num_ch_dec[3])

        self.c2e["3"] = Cube2Equirec(self.cube_h // 8, self.equi_h // 8, self.equi_w // 8)
        self.equi_dec_convs["fusion_3"] = FusionLayer(self.num_ch_enc[2], SE=self.se_in_fusion)
        self.equi_dec_convs["deconv_3"] = ConvBlock(self.num_ch_dec[3] + self.num_ch_enc[2], self.num_ch_dec[3])
        self.equi_dec_convs["upconv_3"] = ConvBlock(self.num_ch_dec[3], self.num_ch_dec[2])

        self.c2e["2"] = Cube2Equirec(self.cube_h // 4, self.equi_h // 4, self.equi_w // 4)
        self.equi_dec_convs["fusion_2"] = FusionLayer(self.num_ch_enc[1], SE=self.se_in_fusion)
        self.equi_dec_convs["deconv_2"] = ConvBlock(self.num_ch_dec[2] + self.num_ch_enc[1], self.num_ch_dec[2])
        self.equi_dec_convs["upconv_2"] = ConvBlock(self.num_ch_dec[2], self.num_ch_dec[1])

        self.c2e["1"] = Cube2Equirec(self.cube_h // 2, self.equi_h // 2, self.equi_w // 2)
        self.equi_dec_convs["fusion_1"] = FusionLayer(self.num_ch_enc[0], SE=self.se_in_fusion)
        self.equi_dec_convs["deconv_1"] = ConvBlock(self.num_ch_dec[1] + self.num_ch_enc[0], self.num_ch_dec[1])
        self.equi_dec_convs["upconv_1"] = ConvBlock(self.num_ch_dec[1], self.num_ch_dec[0])

        self.equi_dec_convs["deconv_0"] = ConvBlock(self.num_ch_dec[0], self.num_ch_dec[0])

        self.equi_dec_convs["depthconv_0"] = Conv3x3(self.num_ch_dec[0], 1)

        self.equi_decoder = nn.ModuleList(list(self.equi_dec_convs.values()))
        self.projectors = nn.ModuleList(list(self.c2e.values()))

        self.sigmoid = nn.Sigmoid()

        self.max_depth = nn.Parameter(torch.tensor(max_depth), requires_grad=False)

    def forward(self, input_equi_image, input_cube_image):

        # euqi image encoding

        if self.num_layers < 18:
            equi_enc_feat0, equi_enc_feat1, equi_enc_feat2, equi_enc_feat3, equi_enc_feat4 \
                = self.equi_encoder(input_equi_image)
        else:
            x = self.equi_encoder.conv1(input_equi_image)
            x = self.equi_encoder.relu(self.equi_encoder.bn1(x))
            equi_enc_feat0 = x

            x = self.equi_encoder.maxpool(x)
            equi_enc_feat1 = self.equi_encoder.layer1(x)
            equi_enc_feat2 = self.equi_encoder.layer2(equi_enc_feat1)
            equi_enc_feat3 = self.equi_encoder.layer3(equi_enc_feat2)
            equi_enc_feat4 = self.equi_encoder.layer4(equi_enc_feat3)

        # cube image encoding
        cube_inputs = torch.cat(torch.split(input_cube_image, self.cube_h, dim=-1), dim=0)

        if self.num_layers < 18:
            cube_enc_feat0, cube_enc_feat1, cube_enc_feat2, cube_enc_feat3, cube_enc_feat4 \
                = self.cube_encoder(cube_inputs)
        else:

            x = self.cube_encoder.conv1(cube_inputs)
            x = self.cube_encoder.relu(self.cube_encoder.bn1(x))
            cube_enc_feat0 = x

            x = self.cube_encoder.maxpool(x)

            cube_enc_feat1 = self.cube_encoder.layer1(x)
            cube_enc_feat2 = self.cube_encoder.layer2(cube_enc_feat1)
            cube_enc_feat3 = self.cube_encoder.layer3(cube_enc_feat2)
            cube_enc_feat4 = self.cube_encoder.layer4(cube_enc_feat3)

        # euqi image decoding fused with cubemap features
        outputs = {}

        cube_enc_feat4 = torch.cat(torch.split(cube_enc_feat4, input_equi_image.shape[0], dim=0), dim=-1)
        c2e_enc_feat4 = self.c2e["5"](cube_enc_feat4)
        fused_feat4 = self.equi_dec_convs["fusion_5"](equi_enc_feat4, c2e_enc_feat4)
        equi_x = upsample(self.equi_dec_convs["upconv_5"](fused_feat4))

        cube_enc_feat3 = torch.cat(torch.split(cube_enc_feat3, input_equi_image.shape[0], dim=0), dim=-1)
        c2e_enc_feat3 = self.c2e["4"](cube_enc_feat3)
        fused_feat3 = self.equi_dec_convs["fusion_4"](equi_enc_feat3, c2e_enc_feat3)
        equi_x = torch.cat([equi_x, fused_feat3], 1)
        equi_x = self.equi_dec_convs["deconv_4"](equi_x)
        equi_x = upsample(self.equi_dec_convs["upconv_4"](equi_x))

        cube_enc_feat2 = torch.cat(torch.split(cube_enc_feat2, input_equi_image.shape[0], dim=0), dim=-1)
        c2e_enc_feat2 = self.c2e["3"](cube_enc_feat2)
        fused_feat2 = self.equi_dec_convs["fusion_3"](equi_enc_feat2, c2e_enc_feat2)
        equi_x = torch.cat([equi_x, fused_feat2], 1)
        equi_x = self.equi_dec_convs["deconv_3"](equi_x)
        equi_x = upsample(self.equi_dec_convs["upconv_3"](equi_x))

        cube_enc_feat1 = torch.cat(torch.split(cube_enc_feat1, input_equi_image.shape[0], dim=0), dim=-1)
        c2e_enc_feat1 = self.c2e["2"](cube_enc_feat1)
        fused_feat1 = self.equi_dec_convs["fusion_2"](equi_enc_feat1, c2e_enc_feat1)
        equi_x = torch.cat([equi_x, fused_feat1], 1)
        equi_x = self.equi_dec_convs["deconv_2"](equi_x)
        equi_x = upsample(self.equi_dec_convs["upconv_2"](equi_x))

        cube_enc_feat0 = torch.cat(torch.split(cube_enc_feat0, input_equi_image.shape[0], dim=0), dim=-1)
        c2e_enc_feat0 = self.c2e["1"](cube_enc_feat0)
        fused_feat0 = self.equi_dec_convs["fusion_1"](equi_enc_feat0, c2e_enc_feat0)
        equi_x = torch.cat([equi_x, fused_feat0], 1)
        equi_x = self.equi_dec_convs["deconv_1"](equi_x)
        equi_x = upsample(self.equi_dec_convs["upconv_1"](equi_x))

        equi_x = self.equi_dec_convs["deconv_0"](equi_x)

        equi_depth = self.equi_dec_convs["depthconv_0"](equi_x)
        outputs["pred_depth"] = self.max_depth * self.sigmoid(equi_depth)

        return outputs


import os
import numpy as np
import kelpie.vasp_output_parser as vasp_output_parser

class DOSError(Exception):
    """Base class to handle error(s) related to density of states."""
    pass

class DOS(object):
    """Base class to store density of states data from a calculation."""

    def __init__(self, vasprunxml_file=None, doscar_file=None):
        """
        :param vasprunxml_file: String with the path to the vasprun.xml file from the calculation.
                                (takes preference over doscar_file)
        :param doscar_file: String with the path to the DOSCAR file from the calculation.
        """
        self._vasprunxml_file = None
        self._vxparser = None
        self.vasprunxml_file = vasprunxml_file

        self._doscar_file = None
        self.doscar_file = doscar_file
        self._doscarparser = None

        self._total_dos_data = self.get_total_dos_data()
        self._fermi_energy = self.get_fermi_energy()
        self._energies = self.get_energies()
        self._total_dos = self.get_total_dos()
        self._total_integrated_dos = self.get_total_integrated_dos()
        self._projected_dos = None

    @property
    def vasprunxml_file(self):
        return self._vasprunxml_file

    @vasprunxml_file.setter
    def vasprunxml_file(self, vasprunxml_file):
        if vasprunxml_file:
            if os.path.isfile(vasprunxml_file):
                self._vasprunxml_file = vasprunxml_file
                self._vxparser = vasp_output_parser.VasprunXMLParser(vasprunxml_file=vasprunxml_file)
            else:
                error_message = 'Cannot find the vasprun.xml file specified: {}'.format(vasprunxml_file)
                raise DOSError(error_message)

    @property
    def vxparser(self):
        return self._vxparser

    @property
    def doscar_file(self):
        return self._doscar_file

    @doscar_file.setter
    def doscar_file(self, doscar_file):
        if doscar_file:
            if os.path.isfile(doscar_file):
                self._doscar_file = doscar_file
            else:
                error_message = 'Cannot find the DOSCAR file specified: {}'.format(doscar_file)
                raise DOSError(error_message)

    @property
    def total_dos_data(self):
        return self._total_dos_data

    def get_total_dos_data(self):
        if self.vxparser is None:
            return
        return self.vxparser.read_total_density_of_states()

    @property
    def fermi_energy(self):
        return self._fermi_energy

    def get_fermi_energy(self):
        if self.vasprunxml_file is not None:
            return self.get_fermi_energy_from_xml()
        elif self.doscar_file is not None:
            return self.get_fermi_energy_from_doscar()

    def get_fermi_energy_from_xml(self):
        return self.vxparser.read_fermi_energy()

    def get_fermi_energy_from_doscar(self):
        raise NotImplementedError

    @property
    def energies(self):
        return self._energies

    def get_energies(self):
        if self.vasprunxml_file is not None:
            return self.get_energies_from_xml()
        elif self.doscar_file is not None:
            return self.get_energies_from_doscar()

    def get_energies_from_xml(self):
        if not self.total_dos_data:
            return
        return np.array([gridpoint[0] - self.fermi_energy for gridpoint in self.total_dos_data['spin_1']])

    def get_energies_from_doscar(self):
        raise NotImplementedError

    @property
    def total_dos(self):
        return self._total_dos

    def get_total_dos(self):
        if self.vasprunxml_file is not None:
            return self.get_total_dos_from_xml()
        elif self.doscar_file is not None:
            return self.get_total_dos_from_doscar()

    def get_total_dos_from_xml(self):
        if not self.total_dos_data:
            return
        total_dos = {}
        for spin in self.total_dos_data:
            total_dos[spin] = np.array([gridpoint[1] for gridpoint in self.total_dos_data[spin]])
        return total_dos

    def get_total_dos_from_doscar(self):
        raise NotImplementedError

    @property
    def total_integrated_dos(self):
        return self._total_integrated_dos

    def get_total_integrated_dos(self):
        if self.vasprunxml_file is not None:
            return self.get_total_integrated_dos_from_xml()
        elif self.doscar_file is not None:
            return self.get_total_integrated_dos_from_doscar()

    def get_total_integrated_dos_from_xml(self):
        if not self.total_dos_data:
            return
        total_intdos = {}
        for spin in self.total_dos_data:
            total_intdos[spin] = np.array([gridpoint[2] for gridpoint in self.total_dos_data[spin]])
        return total_intdos

    def get_total_integrated_dos_from_doscar(self):
        raise NotImplementedError

    @property
    def projected_dos(self):
        return self._projected_dos

    def get_projected_dos(self):
        if self.vasprunxml_file is not None:
            return self.get_projected_dos_from_xml()
        elif self.doscar_file is not None:
            return self.get_projected_dos_from_doscar()

    def get_projected_dos_from_xml(self):
        raise NotImplementedError

    def get_projected_dos_from_doscar(self):
        raise NotImplementedError

    def _check_if_metal(self, spin='spin_1', tol=1e-3):
        if self.energies is None or not self.energies.size:
            return
        efermi_index = np.abs(self.energies).argmin()
        if self.energies[efermi_index] < 0:
            bracket_index = efermi_index + 1
        else:
            bracket_index = efermi_index - 1
        return (self.total_dos[spin][efermi_index] > tol) and (self.total_dos[spin][bracket_index] > tol)

    def check_if_metal(self, tol=1e-3):
        return dict([(spin, self._check_if_metal(spin=spin, tol=tol)) for spin in self.total_dos])

    @property
    def is_metal(self):
        return self.check_if_metal()

    def _find_vbm(self, spin='spin_1', tol=1e-3):
        is_metal = self._check_if_metal(spin=spin, tol=tol)
        if is_metal is None or is_metal:
            return
        efermi_index = np.abs(self.energies).argmin()
        if self.energies[efermi_index] < 0:
            efermi_index += 1
        index = efermi_index
        while self.total_dos[spin][index] < tol:
            index -= 1
        return self.energies[index]

    def find_vbm(self, tol=1e-3):
        return dict([(spin, self._find_vbm(spin=spin, tol=tol)) for spin in self.total_dos])

    @property
    def vbm(self):
        return self.find_vbm()

    def _find_cbm(self, spin='spin_1', tol=1e-3):
        is_metal = self._check_if_metal(spin=spin, tol=tol)
        if is_metal is None or is_metal:
            return
        efermi_index = np.abs(self.energies).argmin()
        if self.energies[efermi_index] > 0:
            efermi_index -= 1
        index = efermi_index
        while self.total_dos[spin][index] < tol:
            index += 1
        return self.energies[index]

    def find_cbm(self, tol=1e-3):
        return dict([(spin, self._find_cbm(spin=spin, tol=tol)) for spin in self.total_dos])

    @property
    def cbm(self):
        return self.find_cbm()

    def _calculate_band_gap(self, spin='spin_1', tol=1e-3):
        is_metal = self._check_if_metal(spin=spin, tol=tol)
        if is_metal is not None:
            return (self.cbm[spin] - self.vbm[spin]) if not is_metal else 0.0

    def calculate_band_gap(self, tol=1e-3):
        return dict([(spin, self._calculate_band_gap(spin=spin, tol=tol)) for spin in self.total_dos])

    @property
    def band_gap(self):
        return self.calculate_band_gap()


# -*- coding: utf-8 -*-
"""
Created on Thu Oct  4 13:30:06 2018

@author: gregz
"""
import time
import numpy as np
import os.path as op
import glob
import warnings

from astropy.io import fits
from utils import biweight_location
from scipy.signal import savgol_filter, medfilt2d
from scipy.interpolate import interp1d, interp2d
from input_utils import setup_logging
from astrometry import Astrometry

dither_pattern = np.array([[0., 0.], [1.27, -0.73], [1.27, 0.73]])
virus_amps = ['LL', 'LU', 'RU', 'RL']
lrs2_amps = [['LL', 'LU'], ['RL', 'RU']]
fplane_file = '/work/03730/gregz/maverick/fplane.txt'
flt_obs = '%07d' % 15
twi_obs = '%07d' % 1
sci_obs = '%07d' % 13
twi_date = '20170205'
sci_date = twi_date
flt_date = twi_date

# FOR LRS2
instrument = 'lrs2'
AMPS = lrs2_amps[0]
dither_pattern = np.zeros((10, 2))

log = setup_logging('panacea_quicklook')

basered = '/work/03730/gregz/maverick'
baseraw = '/work/03946/hetdex/maverick'

twi_path = op.join(basered, 'reductions', twi_date, '%s', '%s%s', 'exp01',
                   '%s', 'multi_*_%s_*_LL.fits')

sci_path = op.join(baseraw, sci_date,  '%s', '%s%s', 'exp%s',
                   '%s', '2*_%sLL*.fits')
flt_path = op.join(baseraw, flt_date,  '%s', '%s%s', 'exp*',
                   '%s', '2*_%sLL*.fits')
sciflt_path = op.join(baseraw, twi_date,  '%s', '%s%s', 'exp*',
                      '%s', '2*_%sLL_twi.fits')
bias_path = op.join(baseraw, twi_date, '%s', '%s%s', 'exp*',
                    '%s', '2*_%sLL_zro.fits')

def get_cal_info(twi_path, amp):
    F = fits.open(glob.glob(twi_path.replace('LL', amp))[0])
    return (np.array(F['ifupos'].data, dtype=float),
            np.array(F['trace'].data, dtype=float),
            np.array(F['wavelength'].data, dtype=float))

def orient_image(image, amp, ampname):
    '''
    Orient the images from blue to red (left to right)
    Fibers are oriented to match configuration files
    '''
    if amp == "LU":
        image[:] = image[::-1, ::-1]
    if amp == "RL":
        image[:] = image[::-1, ::-1]
    if ampname is not None:
        if ampname == 'LR' or ampname == 'UL':
            image[:] = image[:, ::-1]
    return image

def make_avg_spec(wave, spec, binsize=35, per=50):
    ind = np.argsort(wave.ravel())
    T = 1
    for p in wave.shape:
        T *= p
    wchunks = np.array_split(wave.ravel()[ind],
                             T / binsize)
    schunks = np.array_split(spec.ravel()[ind],
                             T / binsize)
    nwave = np.array([np.mean(chunk) for chunk in wchunks])
    nspec = np.array([np.percentile(chunk, per) for chunk in schunks])
    nwave, nind = np.unique(nwave, return_index=True)
    return nwave, nspec[nind]

def base_reduction(filename):
    a = fits.open(filename)
    image = np.array(a[0].data, dtype=float)
    # overscan sub
    overscan_length = 32 * (image.shape[1] / 1064)
    O = biweight_location(image[:, -(overscan_length-2):])
    image[:] = image - O
    # trim image
    image = image[:, :-overscan_length]
    try:
        ampname = a[0].header['AMPNAME']
    except:
        ampname = None
    a = orient_image(image, amp, ampname)
    return a

def get_sciflat_field(flt_path, amp, array_wave, array_trace, common_wave,
                      masterbias, log):
    files = glob.glob(flt_path.replace('LL', amp))
    listflat = []
    array_flt = base_reduction(files[0])

    bigW = np.zeros(array_flt.shape)
    Y, X = np.indices(array_wave.shape)
    YY, XX = np.indices(array_flt.shape)
    for x, at, aw, xx, yy in zip(np.array_split(X, 2, axis=0),
                                 np.array_split(array_trace, 2, axis=0),
                                 np.array_split(array_wave, 2, axis=0),
                                 np.array_split(XX, 2, axis=0),
                                 np.array_split(YY, 2, axis=0)):
        for j in np.arange(at.shape[1]):
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                p0 = np.polyfit(at[:, j], aw[:, j], 7)
            bigW[yy[:, j], j] = np.polyval(p0, yy[:, j])
    listspec = []
    for filename in files:
        log.info('Working on sciflat %s' % filename)
        array_flt = base_reduction(filename) - masterbias
        x = np.arange(array_wave.shape[1])
        spectrum = array_trace * 0.
        for fiber in np.arange(array_wave.shape[0]):
            indl = np.floor(array_trace[fiber]).astype(int)
            indh = np.ceil(array_trace[fiber]).astype(int)
            spectrum[fiber] = array_flt[indl, x] / 2. + array_flt[indh, x] / 2.
        smooth = savgol_filter(spectrum, 315, 1, axis=1)
        avg = biweight_location(smooth, axis=(0,))
        norm = biweight_location(smooth / avg, axis=(1,))
        nw, ns = make_avg_spec(array_wave, spectrum / norm[:, np.newaxis],
                               binsize=41, per=95)
        I = interp1d(nw, ns, kind='linear', fill_value='extrapolate')
        ftf = spectrum * 0.
        for fiber in np.arange(array_wave.shape[0]):
            model = I(array_wave[fiber])
            ftf[fiber] = savgol_filter(spectrum[fiber] / model, 151, 1)
        nw1, ns1 = make_avg_spec(array_wave, spectrum / ftf, binsize=41,
                                 per=95)
        I = interp1d(nw1, ns1, kind='quadratic', fill_value='extrapolate')
        modelimage = I(bigW)
        flat = array_flt / modelimage
        listflat.append(flat)
        listspec.append(I(common_wave))
    flat = np.median(listflat, axis=(0,))
    flat[~np.isfinite(flat)] = 0.0
    flat[flat < 0.0] = 0.0

    return flat, bigW, np.nanmedian(listspec, axis=0)

def safe_division(num, denom, eps=1e-8, fillval=0.0):
    good = np.isfinite(denom) * (np.abs(denom) > eps)
    div = num * 0.
    if num.ndim == denom.ndim:
        div[good] = num[good] / denom[good]
        div[~good] = fillval
    else:
        div[:, good] = num[:, good] / denom[good]
        div[:, ~good] = fillval
    return div

def find_cosmics(Y, E, thresh=8.):
    A = medfilt2d(Y, (5, 1))
    S = safe_division((Y - A), E)
    P = S - medfilt2d(S, (1, 15))
    x, y = np.where(P > thresh)
    xx, yy = ([], [])
    for i in np.arange(-1, 2):
        for j in np.arange(-1, 2):
            sel = ((x + i) >= 0) * ((x + i) < Y.shape[0])
            sel2 = ((y + j) >= 0) * ((y + j) < Y.shape[1])
            sel = sel * sel2
            xx.append((x + i)[sel])
            yy.append((y + j)[sel])
    xx = np.hstack(xx)
    yy = np.hstack(yy)
    inds = np.ravel_multi_index([xx, yy], Y.shape)
    inds = np.unique(inds)
    C = np.zeros(Y.shape, dtype=bool).ravel()
    C[inds] = True
    C = C.reshape(Y.shape)
    log.info('Number of pixels affected by cosmics: %i' % len(x))
    log.info('Fraction of pixels affected by cosmics: %0.5f' %
             (1.*len(inds)/Y.shape[0]/Y.shape[1]))
    return C

def weighted_extraction(image, flat, trace):
    gain = 0.83
    rdnoise = 3.
    I = image * 1.
    I[I < 0.] = 0.
    E = np.sqrt(rdnoise**2 + gain * I) / gain
    E = safe_division(E, flat)
    E[E < 1e-8] = 1e9
    Y = safe_division(image, flat)
    cosmics = find_cosmics(Y, E)
    x = np.arange(trace.shape[1])
    spectrum = 0. * trace
    TT = np.zeros((trace.shape[0], 3, trace.shape[1], 4))
    for fiber in np.arange(trace.shape[0]):
        T = np.zeros((3, trace.shape[1], 4))
        indl = np.floor(trace[fiber]).astype(int)
        flag = False
        for ss, k in enumerate(np.arange(-1, 3)):
            try:
                T[0, :, ss] = Y[indl+k, x]
                T[1, :, ss] = 1. / E[indl+k, x]**2
                T[2, :, ss] = ~cosmics[indl+k, x]
            except:
                v = indl+k
                sel = np.where((v >= 0) * (v < Y.shape[0]))[0]
                T[0, sel, ss] = Y[v[sel], x[sel]]
                T[1, sel, ss] = 1. / E[v[sel], x[sel]]**2
                T[2, sel, ss] = ~cosmics[v[sel], x[sel]]
                flag = True
        if flag:
            if np.mean(indl) > (Y.shape[0]/2.):
                k = 2
            else:
                k = -1
            v = indl+k
            sel = np.where((v >= 0) * (v < len(x)))[0]
            a = np.sum(T[0, sel] * T[1, sel] * T[2, sel], axis=1)
            b = np.sum(T[1, sel] * T[2, sel], axis=1)
            spectrum[fiber, sel] = safe_division(a, b)
        else:
            a = np.sum(T[0] * T[1] * T[2], axis=1)
            b = np.sum(T[1] * T[2], axis=1)
            spectrum[fiber] = safe_division(a, b)
        TT[fiber] = T
    fits.PrimaryHDU(TT).writeto('wtf2.fits', overwrite=True)
    return spectrum

def get_trace_shift(sci_array, flat, array_trace, Yx):
    YM, XM = np.indices(flat.shape)
    inds = np.zeros((3, array_trace.shape[0], array_trace.shape[1]))
    XN = np.round(array_trace)
    inds[0] = XN - 1.
    inds[1] = XN + 0.
    inds[2] = XN + 1.
    inds = np.array(inds, dtype=int)
    Trace = array_trace * 0.
    FlatTrace = array_trace * 0.
    N = YM.max()
    x = np.arange(array_trace.shape[1])
    for i in np.arange(Trace.shape[0]):
        sel = YM[inds[0, i, :], x] >= 0.
        sel = sel * (YM[inds[2, i, :], x] < N)
        xmax = (YM[inds[1, i, sel], x[sel]] -
                (sci_array[inds[2, i, sel], x[sel]] -
                sci_array[inds[0, i, sel], x[sel]]) /
                (2. * (sci_array[inds[2, i, sel], x[sel]] -
                 2. * sci_array[inds[1, i, sel], x[sel]] +
                 sci_array[inds[0, i, sel], x[sel]])))
        Trace[i, sel] = xmax
        xmax = (YM[inds[1, i, sel], x[sel]] - (flat[inds[2, i, sel], x[sel]] -
                flat[inds[0, i, sel], x[sel]]) /
                (2. * (flat[inds[2, i, sel], x[sel]] - 2. *
                 flat[inds[1, i, sel], x[sel]] +
                 flat[inds[0, i, sel], x[sel]])))
        FlatTrace[i, sel] = xmax
    shifts = np.nanmedian(FlatTrace - Trace, axis=1)
    shifts = np.polyval(np.polyfit(np.nanmedian(FlatTrace, axis=1), shifts, 1),
                        Yx)
    fits.HDUList([fits.PrimaryHDU(FlatTrace),
                  fits.ImageHDU(Trace)]).writeto('test_trace.fits',
                                                 overwrite=True)
    return shifts

def subtract_sci(sci_path, flat, array_trace, array_wave, bigW, masterbias):
    files = sorted(glob.glob(sci_path.replace('LL', amp)))
    array_list = []
    for filename in files:
        log.info('Skysubtracting sci %s' % filename)
        array_flt = base_reduction(filename) - masterbias
        array_list.append(array_flt)
    sci_array = np.sum(array_list, axis=0)
    Xx = np.arange(flat.shape[1])
    Yx = np.arange(flat.shape[0])
    I = interp2d(Xx, Yx, flat, kind='cubic', bounds_error=False,
                 fill_value=0.0)
    shifts = get_trace_shift(sci_array, flat, array_trace, Yx)
    flat = I(Xx, Yx + shifts)
    log.info('Found shift for %s of %0.3f' % (files[0], np.median(shifts)))
    array_list = []
    residual = []
    spec_list = []
    for filename in files:
        log.info('Skysubtracting sci %s' % filename)
        array_flt = base_reduction(filename) - masterbias
        array_list.append(array_flt)
        spectrum = weighted_extraction(array_flt, flat, array_trace)
        spectrum[~np.isfinite(spectrum)] = 0.0
        nw, ns = make_avg_spec(array_wave, spectrum, binsize=41)
        ns[~np.isfinite(ns)] = 0.0
        I = interp1d(nw, ns, kind='quadratic', fill_value='extrapolate')
        modelimage = I(bigW)
        residual.append((array_flt - modelimage*flat))
        speclist = []
        for fiber in np.arange(array_wave.shape[0]):
            dlam = np.diff(array_wave[fiber])
            dlam = np.hstack([dlam[0], dlam])
            I = interp1d(array_wave[fiber], spectrum[fiber] / dlam,
                         kind='quadratic', fill_value='extrapolate')
            speclist.append(I(commonwave))
        spec_list.append(np.array(spectrum))
    return np.array(array_list), np.array(residual), np.array(spec_list)

def get_masterbias(zro_path, amp):
    files = glob.glob(zro_path.replace('LL', amp))
    listzro = []
    for filename in files:
        a = base_reduction(filename)
        listzro.append(a)
    return np.median(listzro, axis=0)

# GET ALL VIRUS IFUSLOTS
twilist = glob.glob(twi_path % (instrument, instrument, twi_obs, instrument,
                                '*'))
ifuslots = [op.basename(x).split('_')[2] for x in twilist]
# LRS2-R
ifuslots = ['066']
fiberpos, fiberspec = ([], [])
log.info('Beginning the long haul.')
nexp = len(glob.glob(sci_path % (instrument, instrument, sci_obs, '*',
                                 instrument, ifuslots[0])))
header = fits.open(glob.glob(sci_path % (instrument, instrument, sci_obs, '01',
                                         instrument,
                                         ifuslots[0]))[0])[0].header
PA = float(header['PARANGLE'])
RA = float(header['TRAJRA'])
DEC = float(header['TRAJDEC'])
log.info('Observation at %0.4f %0.4f, PA: %0.3f' % (RA, DEC, PA))
A = Astrometry(RA, DEC, PA, 0., 0., fplane_file=fplane_file)
allflatspec, allspec, allra, alldec, allx, ally, allsub = ([], [], [], [], [],
                                                           [], [])

# Rectified wavelength
commonwave = np.linspace(6450, 8400, 3000)
N = len(ifuslots) * len(virus_amps)
t1 = time.time()
cnt = 0
cnt2 = 0
breakloop = False
for ifuslot in ifuslots:
    for amp in AMPS:
        log.info('Starting on ifuslot, %s, and amp, %s' % (ifuslot, amp))
        twibase = twi_path % (instrument, instrument, twi_obs, instrument,
                              ifuslot)
        amppos, trace, wave = get_cal_info(twibase, amp)
        if wave.ndim == 1:
            log.info('Insufficient cal data for ifuslot, %s, and amp, %s'
                     % (ifuslot, amp))
            continue
        log.info('Getting Masterbias for ifuslot, %s, and amp, %s' %
                 (ifuslot, amp))
        zro_path = bias_path % (instrument, instrument, '00000*', instrument,
                                ifuslot)
        masterbias = get_masterbias(zro_path, amp)
        twibase = sciflt_path % (instrument, instrument, '00000*', instrument,
                                 ifuslot)
        log.info('Getting SciFlat for ifuslot, %s, and amp, %s' %
                 (ifuslot, amp))
        twiflat, bigW, twispec = get_sciflat_field(twibase, amp, wave, trace,
                                                   commonwave, masterbias, log)
        allflatspec.append(twiflat)
        wave = np.array(wave, dtype=float)
        i1 = []
        scifiles = sci_path % (instrument, instrument, sci_obs, '*',
                               instrument, ifuslot)
        images, subimages, spec = subtract_sci(scifiles, twiflat, trace, wave,
                                               bigW, masterbias)
        allsub.append(images)
        allspec.append(spec)
        for i in np.arange(nexp):
            log.info('Getting RA, Dec for exposure, %i,  ifuslot, %s, and amp,'
                     ' %s' % (i+1, ifuslot, amp))
            ra, dec = A.get_ifupos_ra_dec(ifuslot,
                                          amppos[:, 0] + dither_pattern[i, 0],
                                          amppos[:, 1] + dither_pattern[i, 1])
            allra.append(ra)
            alldec.append(dec)
            allx.append(A.fplane.by_ifuslot(ifuslot).y + amppos[:, 0] +
                        dither_pattern[i, 0])
            ally.append(A.fplane.by_ifuslot(ifuslot).x + amppos[:, 1] +
                        dither_pattern[i, 1])

        t2 = time.time()
        cnt += 1
        time_per_amp = (t2 - t1) / cnt
        remaining_amps = (N - cnt)
        log.info('Time remaining: %0.2f' % (time_per_amp * remaining_amps))
    if breakloop:
        break

fitslist = [fits.PrimaryHDU(np.vstack(allspec)),
            fits.ImageHDU(np.array(allflatspec)),
            fits.ImageHDU(commonwave),
            fits.ImageHDU(np.array(allra)),
            fits.ImageHDU(np.array(alldec)),
            fits.ImageHDU(np.array(allx)),
            fits.ImageHDU(np.array(ally))]
fits.HDUList(fitslist).writeto('test_big.fits', overwrite=True)
flist1 = []
alls = np.vstack(allsub)
for j, resi in enumerate(alls):
    if j == 0:
        func = fits.PrimaryHDU
    else:
        func = fits.ImageHDU
    flist1.append(func(resi))
fits.HDUList(flist1).writeto('test_sub.fits', overwrite=True)


# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import logging
import os
import json_tricks as json
from collections import OrderedDict

import numpy as np

from lib.dataset.JointsDataset import JointsDataset

import imagesize
import cv2
from lib.nanit_utils.utils_functions import transform_box_xywh, gen_bbox_from_xywh_list,\
    get_rotation_matrix_head_up, gen_mpii_joints_mapping_dict, gen_oks_sigmas_dict, gen_mpii_flip_joints_list

logger = logging.getLogger(__name__)

class NanitMPIIDataset(JointsDataset):
    def __init__(self, cfg, root, image_set, is_train, transform=None):
        super().__init__(cfg, root, image_set, is_train, transform, gen_mpii_joints_mapping_dict())

        self.joints_mapping = gen_mpii_joints_mapping_dict()
        _joints_with_zero_weight = self.joints_weight == 0
        self.joints_with_zero_weight = _joints_with_zero_weight.reshape(1, -1)[0]

        self.num_joints = len(self.joints_mapping)
        self.flip_pairs = gen_mpii_flip_joints_list()

        self.upper_body_ids = (self.joints_mapping['center'], self.joints_mapping['neck'], self.joints_mapping['head_top'],
                               self.joints_mapping['wrist_right'], self.joints_mapping['elbow_right'], self.joints_mapping['shoulder_right'],
                               self.joints_mapping['shoulder_left'], self.joints_mapping['elbow_left'], self.joints_mapping['wrist_left'])
        self.lower_body_ids = (self.joints_mapping['ankle_right'], self.joints_mapping['knee_right'], self.joints_mapping['hip_right'],
                               self.joints_mapping['hip_left'], self.joints_mapping['knee_left'], self.joints_mapping['ankle_left'],
                               self.joints_mapping['pelvis'], self.joints_mapping['heel_right'], self.joints_mapping['heel_left'])

        self.db = self._get_db()

        if is_train and cfg.DATASET.SELECT_DATA:
            self.db = self.select_data(self.db)

        # Validation properties
        self.gt_valid_path = None
        self.gt_valid_data = None
        self.jnt_visible = None
        self.pos_gt_src = None
        self.headboxes_src = None
        self.oks_sigmas = None
        self.joints_labels_wandb = {v + 1: k for k, v in self.joints_mapping.items()}
        self.pckh_sc_bias = 0.6      # proportion of head box size for pckh evaluation
        self.pckh_threshold = 0.5    # threshold for pckh evaluation

        if 'val' in self.image_set:
            self.handle_validation_dataset(cfg)

        logger.info('=> load {} samples'.format(len(self.db)))

    def handle_validation_dataset(self, cfg):
        self.gt_valid_path = os.path.join(cfg.DATASET.ROOT, 'annot', '{}.json'.format(cfg.DATASET.TEST_SET))
        with open(self.gt_valid_path) as f:
            self.gt_valid_data = json.load(f)
        self.jnt_visible = np.zeros((self.num_joints, len(self.db)))
        self.pos_gt_src = np.zeros((self.num_joints, 2, len(self.db)))
        self.headboxes_src = np.zeros((2, 2, len(self.db)))
        images_dir = os.path.join(cfg.DATASET.ROOT, 'images')
        for i, d in enumerate(self.gt_valid_data):
            joints_arr = np.array(d['joints'])
            joints_vis_arr = np.array(d['joints_vis'])

            # set joints with zero weight as not visible
            joints_arr[self.joints_with_zero_weight, 0:2] = -1
            joints_vis_arr[self.joints_with_zero_weight] = 0
            self.jnt_visible[:, i] = joints_vis_arr

            if self.rotate_head_up_flag:        # align GT in-case of rotation heads-up
                img_path = os.path.join(images_dir, d['image'])
                w, h = imagesize.get(img_path)
                baby_bbox = gen_bbox_from_xywh_list(d['baby'])
                head_bbox = gen_bbox_from_xywh_list(d['headbbox'])
                rotation_matrix = get_rotation_matrix_head_up(baby_bbox, head_bbox, h, w)

                d['headbbox'] = transform_box_xywh(d['headbbox'], rotation_matrix)                 # keep metadata updated
                if joints_vis_arr.any():
                    joints_for_transformation = joints_arr[joints_vis_arr == 1].reshape(-1, 1, 2)
                    joints_transformed = cv2.transform(joints_for_transformation, rotation_matrix)
                    joints_arr[joints_vis_arr == 1] = joints_transformed[:, 0, :]

            # keep metadata updated
            d['joints_vis'] = list(joints_vis_arr)
            d['joints'] = [list(p) for p in joints_arr]

            self.pos_gt_src[:, :, i] = joints_arr
            self.headboxes_src[:, :, i] = np.array(d['headbbox']).reshape((2, 2))

        self._generate_oks_sigmas()

    def _generate_oks_sigmas(self):
        oks_sigmas_dict = gen_oks_sigmas_dict()
        self.oks_sigmas = np.ones((len(oks_sigmas_dict), 1))
        for k, v in oks_sigmas_dict.items():
            self.oks_sigmas[self.joints_mapping[k]] = v

    def _get_db(self):
        # create train/val split
        file_name = os.path.join(
            self.root, 'annot', self.image_set+'.json'
        )
        with open(file_name) as anno_file:
            anno = json.load(anno_file)

        gt_db = []
        for a in anno:
            image_name = a['image']

            center = np.array(a['center'], dtype=np.float)
            scale = np.array(a['scale'], dtype=np.float)

            # Adjust center/scale slightly to avoid cropping limbs
            if center[0] != -1:
                center[1] = center[1] + 15 * scale[1]
                scale = scale * 1.25

            joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
            joints_3d_vis = np.zeros((self.num_joints,  3), dtype=np.float)
            if self.image_set != 'test':
                joints = np.array(a['joints'])
                joints_vis = np.array(a['joints_vis'])
                assert len(joints) == self.num_joints, \
                    'joint num diff: {} vs {}'.format(len(joints),
                                                      self.num_joints)

                # Set joints with zero weight as not visible
                joints[self.joints_with_zero_weight, 0:2] = -1
                joints_vis[self.joints_with_zero_weight] = 0

                joints_3d[:, 0:2] = joints[:, 0:2]
                joints_3d_vis[:, 0] = joints_vis[:]
                joints_3d_vis[:, 1] = joints_vis[:]

            gt_db.append(
                {
                    'image': os.path.join(self.root, 'images', image_name),
                    'center': center,
                    'scale': scale,
                    'joints_3d': joints_3d,
                    'joints_3d_vis': joints_3d_vis,
                    'filename': '',
                    'imgnum': 0,
                    'headbbox': a['headbbox'],
                    'baby': a['baby']
                }
            )

        return gt_db

    def evaluate(self, cfg, preds, output_dir, *args, **kwargs):
        preds = preds[:, :, 0:2]    # drop the 3 coordinate, use (x, y)

        if output_dir:
            pred_file = os.path.join(output_dir, 'pred.json')
            with open(pred_file, 'w') as f:
                json.dump(preds.tolist(), f)

        if 'test' in cfg.DATASET.TEST_SET:
            return {'Null': 0.0}, 0.0

        name_value = self.pckh_and_oks_evaluate(preds)

        return name_value, name_value['Mean@0.1']       # Best Model indicator set to Mean@0.1

    def pckh_and_oks_evaluate(self, preds):
        pos_pred_src = np.transpose(preds, [1, 2, 0])
        head = self.joints_mapping['head_top']
        neck = self.joints_mapping['neck']
        lsho = self.joints_mapping['shoulder_left']
        lelb = self.joints_mapping['elbow_left']
        lwri = self.joints_mapping['wrist_left']
        lhip = self.joints_mapping['hip_left']
        lkne = self.joints_mapping['knee_left']
        lank = self.joints_mapping['ankle_left']
        rsho = self.joints_mapping['shoulder_right']
        relb = self.joints_mapping['elbow_right']
        rwri = self.joints_mapping['wrist_right']
        rkne = self.joints_mapping['hip_right']
        rank = self.joints_mapping['knee_right']
        rhip = self.joints_mapping['ankle_right']
        rheel = self.joints_mapping['heel_right']
        lheel = self.joints_mapping['heel_left']

        pos_gt_src = self.pos_gt_src
        headboxes_src = self.headboxes_src
        jnt_visible = self.jnt_visible

        uv_error = pos_pred_src - pos_gt_src
        uv_err = np.linalg.norm(uv_error, axis=1)
        headsizes = headboxes_src[1, :, :]  # headboxes_src are x,y,w,h, shape 2x2xN. sizes (w,h) are headboxes_src[1, :, :], shape 2xN
        headsizes = np.linalg.norm(headsizes, axis=0)  # shape 1XN
        headsizes *= self.pckh_sc_bias
        scale = np.multiply(headsizes, np.ones((len(uv_err), 1)))
        scaled_uv_err = np.divide(uv_err, scale)
        scaled_uv_err = np.multiply(scaled_uv_err, jnt_visible)
        jnt_count = np.sum(jnt_visible, axis=1)
        less_than_threshold = np.multiply((scaled_uv_err <= self.pckh_threshold),
                                          jnt_visible)
        PCKh = np.divide(100. * np.sum(less_than_threshold, axis=1), jnt_count)

        # save
        rng = np.arange(0, 0.5 + 0.01, 0.01)
        pckAll = np.zeros((len(rng), self.num_joints))
        for r in range(len(rng)):
            threshold = rng[r]
            less_than_threshold = np.multiply(scaled_uv_err <= threshold, jnt_visible)
            pckAll[r, :] = np.divide(100. * np.sum(less_than_threshold, axis=1), jnt_count)

        oks = np.zeros((len(self.joints_mapping)))
        scale_oks = 2.0 * scale * self.oks_sigmas
        for joint_index in self.joints_mapping.values():
            scaled_uv_err_oks = np.power(np.divide(uv_err, scale_oks), 2)[joint_index, jnt_visible[joint_index] == 1]
            oks[joint_index] = 100. * np.sum(np.exp(-1.0 * scaled_uv_err_oks)) / len(scaled_uv_err_oks)

        PCKh = np.ma.array(PCKh, mask=False)
        PCKh.mask[self.joints_mapping['pelvis']] = True
        PCKh.mask[self.joints_mapping['center']] = True

        jnt_count = np.ma.array(jnt_count, mask=False)
        jnt_count.mask[self.joints_mapping['pelvis']] = True
        jnt_count.mask[self.joints_mapping['center']] = True

        jnt_ratio = jnt_count / np.sum(jnt_count).astype(np.float64)

        rng_index_for_0_1 = 10
        sms_joints = [head, neck, lhip, rhip, lkne, rkne, lheel, rheel]
        sms_joints_pckh_0_1 = [pckAll[rng_index_for_0_1, j] for j in sms_joints]
        sms_joints_oks = [oks[j] for j in sms_joints]

        name_value = [
            ('Head', PCKh[head]),
            ('Neck', PCKh[neck]),
            ('Shoulder', 0.5 * (PCKh[lsho] + PCKh[rsho])),
            ('Elbow', 0.5 * (PCKh[lelb] + PCKh[relb])),
            ('Wrist', 0.5 * (PCKh[lwri] + PCKh[rwri])),
            ('Hip', 0.5 * (PCKh[lhip] + PCKh[rhip])),
            ('Knee', 0.5 * (PCKh[lkne] + PCKh[rkne])),
            ('Ankle', 0.5 * (PCKh[lank] + PCKh[rank])),
            ('Heel', 0.5 * (PCKh[lheel] + PCKh[rheel])),
            ('Mean', np.sum(PCKh * jnt_ratio)),

            ('Head@0.1', pckAll[rng_index_for_0_1, head]),
            ('Neck@0.1', pckAll[rng_index_for_0_1, neck]),
            ('Shoulder@0.1', 0.5 * (pckAll[rng_index_for_0_1, lsho] + pckAll[rng_index_for_0_1, rsho])),
            ('Elbow@0.1', 0.5 * (pckAll[rng_index_for_0_1, lelb] + pckAll[rng_index_for_0_1, relb])),
            ('Wrist@0.1', 0.5 * (pckAll[rng_index_for_0_1, lwri] + pckAll[rng_index_for_0_1, rwri])),
            ('Hip@0.1', 0.5 * (pckAll[rng_index_for_0_1, lhip] + pckAll[rng_index_for_0_1, rhip])),
            ('Knee@0.1', 0.5 * (pckAll[rng_index_for_0_1, lkne] + pckAll[rng_index_for_0_1, rkne])),
            ('Ankle@0.1', 0.5 * (pckAll[rng_index_for_0_1, lank] + pckAll[rng_index_for_0_1, rank])),
            ('Heel@0.1', 0.5 * (pckAll[rng_index_for_0_1, lheel] + pckAll[rng_index_for_0_1, rheel])),
            ('Mean@0.1', np.sum(pckAll[rng_index_for_0_1, :] * jnt_ratio)),
            ('MeanSmS@0.1', np.sum(sms_joints_pckh_0_1) / len(sms_joints)),

            ('Head_OKS', oks[head]),
            ('Neck_OKS', oks[neck]),
            ('Shoulder_OKS', 0.5 * (oks[lsho] + oks[rsho])),
            ('Elbow_OKS', 0.5 * (oks[lelb] + oks[relb])),
            ('Wrist_OKS', 0.5 * (oks[lwri] + oks[rwri])),
            ('Hip_OKS', 0.5 * (oks[lhip] + oks[rhip])),
            ('Knee_OKS', 0.5 * (oks[lkne] + oks[rkne])),
            ('Ankle_OKS', 0.5 * (oks[lank] + oks[rank])),
            ('Heel_OKS', 0.5 * (oks[lheel] + oks[rheel])),
            ('Mean_OKS', np.sum(oks * jnt_ratio)),
            ('MeanSmS_OKS', np.sum(sms_joints_oks) / len(sms_joints)),

        ]
        name_value = OrderedDict(name_value)

        return name_value


"""
Retrain the YOLO model for your own dataset.
"""
import os, time, random, argparse
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

import numpy as np
import h5py
import tensorflow.keras.backend as K
from tensorflow.distribute import MirroredStrategy as multi_gpu_model
from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, LearningRateScheduler, EarlyStopping, TerminateOnNaN, LambdaCallback
from tensorflow_model_optimization.sparsity import keras as sparsity

from yolo5.model import get_yolo5_train_model
from yolo3.model import get_yolo3_train_model
from yolo3.data import yolo3_data_generator_wrapper, Yolo3DataGenerator
from yolo3.loss import yolo3_loss
from yolo2.model import get_yolo2_train_model
from yolo2.data import yolo2_data_generator_wrapper, Yolo2DataGenerator
from common.utils import get_classes, get_anchors, get_dataset, optimize_tf_gpu
from common.model_utils import get_optimizer
from common.callbacks import EvalCallBack, DatasetShuffleCallBack

# Try to enable Auto Mixed Precision on TF 2.0
os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'
os.environ['TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_IGNORE_PERFORMANCE'] = '1'

import tensorflow as tf
optimize_tf_gpu(tf, K)

# PARAMETERS
val_split = 0.3
batch_size = 16
model_image_size = (224, 224)
enhance_augment = None #/ "mosaic"
rescale_interval = 10
epochs = 5

# GET DATA
anchors = get_anchors("configs/yolo3_anchors.txt")
classes = get_classes("configs/epic_kitchens_classes.txt")
num_classes = len(classes)
num_anchors = len(anchors)
num_feature_layers = num_anchors // 3
model_type = "yolo3_vgg16"
weights_path = None
# weights = h5py.File("./my_model.h5", 'r')
dataset = get_dataset("epic_kitchens_train_data.txt")
num_val = int(len(dataset)*val_split)
num_train = len(dataset) - num_val
rescale_interval = -1

# SET UP LOGGING
logging = TensorBoard(log_dir="logs/001", histogram_freq=0, write_graph=False, write_grads=False, write_images=False, update_freq='batch')
checkpoint = ModelCheckpoint(os.path.join(".", 'currweights.h5'),
    monitor='val_loss',
    mode='min',
    verbose=1,
    save_weights_only=False,
    save_best_only=True,
    period=1)
reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, mode='min', patience=10, verbose=1, cooldown=0, min_lr=1e-10)
early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=50, verbose=1, mode='min')
terminate_on_nan = TerminateOnNaN()
shuffle_callback = DatasetShuffleCallBack(dataset)
callbacks=[logging, checkpoint, reduce_lr, early_stopping, terminate_on_nan, shuffle_callback]

# GET TRAINING MODELS AND DATA
train_model = get_yolo3_train_model(model_type, anchors, num_classes, model_pruning=False, weights_path=None)
optimizer = get_optimizer('adam', 0.001, decay_type=None)
train_data = yolo3_data_generator_wrapper(dataset[:num_train], batch_size, model_image_size, anchors, num_classes, enhance_augment, rescale_interval, multi_anchor_assign=False)
validation_data = yolo3_data_generator_wrapper(dataset[num_train:], batch_size, model_image_size, anchors, num_classes, multi_anchor_assign=False)

# # TRAIN
# train_model.fit(train_data,
#                 steps_per_epoch=max(1, num_train//batch_size),
#                 validation_data=validation_data,
#                 validation_steps=max(1,num_val//batch_size),
#                 epochs=epochs,
#                 initial_epoch=1,
#                 workers=1,
#                 use_multiprocessing=False,
#                 max_queue_size=10,
#                 callbacks=callbacks)

# # quit()
# time.sleep(2)

# TRAIN
for i in range(len(train_model.layers)):
    train_model.layers[i].trainable = True
train_model.compile(optimizer=optimizer, loss=train_model.loss)
print('Train on {} samples, val on {} samples, with batch size {}, input_shape {}.'.format(num_train, num_val, batch_size, model_image_size))
train_model.fit(train_data,
                steps_per_epoch=max(1, num_train//batch_size),
                validation_data=validation_data,
                validation_steps=max(1,num_val//batch_size),
                epochs=epochs,
                initial_epoch=1,
                workers=1,
                use_multiprocessing=False,
                max_queue_size=10,
                callbacks=callbacks)
train_model.save(os.path.join(".", 'my_model.h5'))
train_model.summary()
# print(yolo3_model)


from PySide2 import QtCore, QtGui
from PySide2.QtCore import Signal, Slot
from PIL import Image, ImageDraw, ImageFont
from PIL.ImageQt import ImageQt
import core
import numpy
import subprocess as sp
import sys

class Worker(QtCore.QObject):

  videoCreated = Signal()
  progressBarUpdate = Signal(int)
  progressBarSetText = Signal(str)

  def __init__(self, parent=None):
    QtCore.QObject.__init__(self)
    parent.videoTask.connect(self.createVideo)
    self.core = core.Core()

  @Slot(str, str, QtGui.QFont, int, int, int, int, tuple, tuple, str, str)
  def createVideo(self, backgroundImage, titleText, titleFont, fontSize, alignment,\
                    xOffset, yOffset,  textColor, visColor, inputFile, outputFile):
    # print('worker thread id: {}'.format(QtCore.QThread.currentThreadId()))
    def getBackgroundAtIndex(i):
        return self.core.drawBaseImage(
            backgroundFrames[i],
            titleText,
            titleFont,
            fontSize,
            alignment,
            xOffset,
            yOffset,
            textColor,
            visColor)

    progressBarValue = 0
    self.progressBarUpdate.emit(progressBarValue)
    self.progressBarSetText.emit('Loading background image…')

    backgroundFrames = self.core.parseBaseImage(backgroundImage)
    if len(backgroundFrames) < 2:
        # the base image is not a video so we can draw it now
        imBackground = getBackgroundAtIndex(0)
    else:
        # base images will be drawn while drawing the audio bars
        imBackground = None

    self.progressBarSetText.emit('Loading audio file…')
    completeAudioArray = self.core.readAudioFile(inputFile)

    # test if user has libfdk_aac
    encoders = sp.check_output(self.core.FFMPEG_BIN + " -encoders -hide_banner", shell=True)
    if b'libfdk_aac' in encoders:
      acodec = 'libfdk_aac'
    else:
      acodec = 'aac'

    ffmpegCommand = [ self.core.FFMPEG_BIN,
       '-y', # (optional) means overwrite the output file if it already exists.
       '-f', 'rawvideo',
       '-vcodec', 'rawvideo',
       '-s', '1280x720', # size of one frame
       '-pix_fmt', 'rgb24',
       '-r', '30', # frames per second
       '-i', '-', # The input comes from a pipe
       '-i', inputFile,
       '-acodec', acodec, # output audio codec
       '-b:a', "192k",
       '-vcodec', "libx264",
       '-pix_fmt', "yuv420p",
       '-preset', "medium",
       '-f', "mp4"]

    if acodec == 'aac':
      ffmpegCommand.append('-strict')
      ffmpegCommand.append('-2')

    ffmpegCommand.append(outputFile)

    out_pipe = sp.Popen(ffmpegCommand,
        stdin=sp.PIPE,stdout=sys.stdout, stderr=sys.stdout)

    smoothConstantDown = 0.08
    smoothConstantUp = 0.8
    lastSpectrum = None
    sampleSize = 1470

    numpy.seterr(divide='ignore')
    bgI = 0
    for i in range(0, len(completeAudioArray), sampleSize):
      # create video for output
      lastSpectrum = self.core.transformData(
        i,
        completeAudioArray,
        sampleSize,
        smoothConstantDown,
        smoothConstantUp,
        lastSpectrum)
      if imBackground != None:
        im = self.core.drawBars(lastSpectrum, imBackground, visColor)
      else:
        im = self.core.drawBars(lastSpectrum, getBackgroundAtIndex(bgI), visColor)
        if bgI < len(backgroundFrames)-1:
            bgI += 1

      # write to out_pipe
      try:
        out_pipe.stdin.write(im.tobytes())
      finally:
        True

      # increase progress bar value
      if progressBarValue + 1 <= (i / len(completeAudioArray)) * 100:
        progressBarValue = numpy.floor((i / len(completeAudioArray)) * 100)
        self.progressBarUpdate.emit(progressBarValue)
        self.progressBarSetText.emit('%s%%' % str(int(progressBarValue)))

    numpy.seterr(all='print')

    out_pipe.stdin.close()
    if out_pipe.stderr is not None:
      print(out_pipe.stderr.read())
      out_pipe.stderr.close()
    # out_pipe.terminate() # don't terminate ffmpeg too early
    out_pipe.wait()
    print("Video file created")
    self.core.deleteTempDir()
    self.progressBarUpdate.emit(100)
    self.progressBarSetText.emit('100%')
    self.videoCreated.emit()