from easydict import EasyDict as edict
import math
import numpy as np
# make training faster
# our RAM is 256G
# mount -t tmpfs -o size=140G  tmpfs /train_tmp

config = edict()
config.loss = "cosface"
config.network = "r50"
config.resume = False
config.output = "ms1mv3_arcface_r50"

config.dataset = "ms1m-retinaface-t1"
config.embedding_size = 512
config.sample_rate = 1
config.fp16 = True
config.momentum = 0.9
config.weight_decay = 5e-4
config.batch_size = 128
config.lr = 0.1  # batch size is 512
config.model_load_dir = ''

config.val_batch_size = 10

config.node_ips = ["192.168.1.13"]
config.num_nodes = 1
config.device_num_per_node = 1
config.model_parallel = 1
config.partial_fc = 0

config.use_synthetic_data = False

config.fc_type = "FC"
config.nccl_fusion_threshold_mb = 16
config.nccl_fusion_max_ops = 64
config.val_dataset_dir = "/train_tmp/glint360k/val"

config.part_name_prefix = "part-"
config.part_name_suffix_length = 5
config.train_data_part_num = 16
config.shuffle = True

if config.dataset == "emore":
    config.ofrecord_path = "/train_tmp/faces_emore"
    config.num_classes = 85742
    config.num_image = 5822653
    config.num_epoch = 16
    config.warmup_epoch = -1
    config.decay_epoch = [8, 14, ]
    config.val_targets = ["lfw", ]
    config.train_data_part_num = 32

elif config.dataset == "ms1m-retinaface-t1":
    config.ofrecord_path = "/dev/shm/ms1m-retinaface-t1/ofrecord"
    config.num_classes = 93432
    config.num_image = 5179510
    config.num_epoch = 25
    config.warmup_epoch = -1
    config.decay_epoch = [11, 17, 22]
    config.val_targets = ["lfw", "cfp_fp", "agedb_30"]
    config.train_data_part_num = 32

elif config.dataset == "glint360k":
    config.ofrecord_path = "/train_tmp/glint360k"
    config.num_classes = 360232
    config.num_image = 17091657
    config.num_epoch = 20
    config.warmup_epoch = -1
    config.decay_epoch = [8, 12, 15, 18]
    config.val_targets = ["lfw", "cfp_fp", "agedb_30"]

elif config.dataset == "webface":
    config.ofrecord_path = "/train_tmp/faces_webface_112x112"
    config.num_classes = 10572
    config.num_image = "forget"
    config.num_epoch = 34
    config.warmup_epoch = -1
    config.decay_epoch = [20, 28, 32]
    config.val_targets = ["lfw", "cfp_fp", "agedb_30"]


import os
from collections import OrderedDict
import numpy as np
from rdkit import Chem
import networkx as nx
import torch
from torch_geometric import data as DATA

# nomarlize
def dic_normalize(dic):
    # print(dic)
    max_value = dic[max(dic, key=dic.get)]
    min_value = dic[min(dic, key=dic.get)]
    # print(max_value)
    interval = float(max_value) - float(min_value)
    for key in dic.keys():
        dic[key] = (dic[key] - min_value) / interval
    dic['X'] = (max_value + min_value) / 2.0
    return dic

pro_res_table = ['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y', 'X']

pro_res_aliphatic_table = ['A', 'I', 'L', 'M', 'V']
pro_res_aromatic_table = ['F', 'W', 'Y']
pro_res_polar_neutral_table = ['C', 'N', 'Q', 'S', 'T']
pro_res_acidic_charged_table = ['D', 'E']
pro_res_basic_charged_table = ['H', 'K', 'R']

res_weight_table = {'A': 71.08, 'C': 103.15, 'D': 115.09, 'E': 129.12, 'F': 147.18, 'G': 57.05, 'H': 137.14,
                    'I': 113.16, 'K': 128.18, 'L': 113.16, 'M': 131.20, 'N': 114.11, 'P': 97.12, 'Q': 128.13,
                    'R': 156.19, 'S': 87.08, 'T': 101.11, 'V': 99.13, 'W': 186.22, 'Y': 163.18}

res_pka_table = {'A': 2.34, 'C': 1.96, 'D': 1.88, 'E': 2.19, 'F': 1.83, 'G': 2.34, 'H': 1.82, 'I': 2.36,
                 'K': 2.18, 'L': 2.36, 'M': 2.28, 'N': 2.02, 'P': 1.99, 'Q': 2.17, 'R': 2.17, 'S': 2.21,
                 'T': 2.09, 'V': 2.32, 'W': 2.83, 'Y': 2.32}

res_pkb_table = {'A': 9.69, 'C': 10.28, 'D': 9.60, 'E': 9.67, 'F': 9.13, 'G': 9.60, 'H': 9.17,
                 'I': 9.60, 'K': 8.95, 'L': 9.60, 'M': 9.21, 'N': 8.80, 'P': 10.60, 'Q': 9.13,
                 'R': 9.04, 'S': 9.15, 'T': 9.10, 'V': 9.62, 'W': 9.39, 'Y': 9.62}

res_pkx_table = {'A': 0.00, 'C': 8.18, 'D': 3.65, 'E': 4.25, 'F': 0.00, 'G': 0, 'H': 6.00,
                 'I': 0.00, 'K': 10.53, 'L': 0.00, 'M': 0.00, 'N': 0.00, 'P': 0.00, 'Q': 0.00,
                 'R': 12.48, 'S': 0.00, 'T': 0.00, 'V': 0.00, 'W': 0.00, 'Y': 0.00}

res_pl_table = {'A': 6.00, 'C': 5.07, 'D': 2.77, 'E': 3.22, 'F': 5.48, 'G': 5.97, 'H': 7.59,
                'I': 6.02, 'K': 9.74, 'L': 5.98, 'M': 5.74, 'N': 5.41, 'P': 6.30, 'Q': 5.65,
                'R': 10.76, 'S': 5.68, 'T': 5.60, 'V': 5.96, 'W': 5.89, 'Y': 5.96}

res_hydrophobic_ph2_table = {'A': 47, 'C': 52, 'D': -18, 'E': 8, 'F': 92, 'G': 0, 'H': -42, 'I': 100,
                             'K': -37, 'L': 100, 'M': 74, 'N': -41, 'P': -46, 'Q': -18, 'R': -26, 'S': -7,
                             'T': 13, 'V': 79, 'W': 84, 'Y': 49}

res_hydrophobic_ph7_table = {'A': 41, 'C': 49, 'D': -55, 'E': -31, 'F': 100, 'G': 0, 'H': 8, 'I': 99,
                             'K': -23, 'L': 97, 'M': 74, 'N': -28, 'P': -46, 'Q': -10, 'R': -14, 'S': -5,
                             'T': 13, 'V': 76, 'W': 97, 'Y': 63}

res_weight_table = dic_normalize(res_weight_table)
res_pka_table = dic_normalize(res_pka_table)
res_pkb_table = dic_normalize(res_pkb_table)
res_pkx_table = dic_normalize(res_pkx_table)
res_pl_table = dic_normalize(res_pl_table)
res_hydrophobic_ph2_table = dic_normalize(res_hydrophobic_ph2_table)
res_hydrophobic_ph7_table = dic_normalize(res_hydrophobic_ph7_table)

def residue_features(residue):
    res_property1 = [1 if residue in pro_res_aliphatic_table else 0, 1 if residue in pro_res_aromatic_table else 0,
                     1 if residue in pro_res_polar_neutral_table else 0,
                     1 if residue in pro_res_acidic_charged_table else 0,
                     1 if residue in pro_res_basic_charged_table else 0]
    res_property2 = [res_weight_table[residue], res_pka_table[residue], res_pkb_table[residue], res_pkx_table[residue],
                     res_pl_table[residue], res_hydrophobic_ph2_table[residue], res_hydrophobic_ph7_table[residue]]
    # print(np.array(res_property1 + res_property2).shape)
    return np.array(res_property1 + res_property2)

# mol atom feature for mol graph
def atom_features(atom):
    # 44 +11 +11 +11 +1
    return np.array(one_of_k_encoding_unk(atom.GetSymbol(),
                                          ['C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na', 'Ca', 'Fe', 'As',
                                           'Al', 'I', 'B', 'V', 'K', 'Tl', 'Yb', 'Sb', 'Sn', 'Ag', 'Pd', 'Co', 'Se',
                                           'Ti', 'Zn', 'H', 'Li', 'Ge', 'Cu', 'Au', 'Ni', 'Cd', 'In', 'Mn', 'Zr', 'Cr',
                                           'Pt', 'Hg', 'Pb', 'X']) +
                    one_of_k_encoding(atom.GetDegree(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +
                    one_of_k_encoding_unk(atom.GetTotalNumHs(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +
                    one_of_k_encoding_unk(atom.GetImplicitValence(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +
                    [atom.GetIsAromatic()])

# one ont encoding
def one_of_k_encoding(x, allowable_set):
    if x not in allowable_set:
        # print(x)
        raise Exception('input {0} not in allowable set{1}:'.format(x, allowable_set))
    return list(map(lambda s: x == s, allowable_set))

def one_of_k_encoding_unk(x, allowable_set):
    '''Maps inputs not in the allowable set to the last element.'''
    if x not in allowable_set:
        x = allowable_set[-1]
    return list(map(lambda s: x == s, allowable_set))

# mol smile to mol graph edge index
def smile_to_graph(smile):
    mol = Chem.MolFromSmiles(smile)

    c_size = mol.GetNumAtoms()

    features = []
    for atom in mol.GetAtoms():
        feature = atom_features(atom)
        features.append(feature / sum(feature))

    edges = []
    for bond in mol.GetBonds():
        edges.append([bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()])
    g = nx.Graph(edges).to_directed()
    edge_index = []
    mol_adj = np.zeros((c_size, c_size))
    for e1, e2 in g.edges:
        mol_adj[e1, e2] = 1
        # edge_index.append([e1, e2])
    mol_adj += np.matrix(np.eye(mol_adj.shape[0]))
    index_row, index_col = np.where(mol_adj >= 0.5)
    for i, j in zip(index_row, index_col):
        edge_index.append([i, j])
    # print('smile_to_graph')
    return c_size, features, edge_index

# target feature for target graph
def PSSM_calculation(aln_file, pro_seq):
    pfm_mat = np.zeros((len(pro_res_table), len(pro_seq)))
    with open(aln_file, 'r') as f:
        line_count = len(f.readlines())
        for line in f.readlines():
            if len(line) != len(pro_seq):
                print('error', len(line), len(pro_seq))
                continue
            count = 0
            for res in line:
                if res not in pro_res_table:
                    count += 1
                    continue
                pfm_mat[pro_res_table.index(res), count] += 1
                count += 1
    # ppm_mat = pfm_mat / float(line_count)
    pseudocount = 0.8
    ppm_mat = (pfm_mat + pseudocount / 4) / (float(line_count) + pseudocount)
    pssm_mat = ppm_mat
    # k = float(len(pro_res_table))
    # pwm_mat = np.log2(ppm_mat / (1.0 / k))
    # pssm_mat = pwm_mat
    # print(pssm_mat)
    return pssm_mat

def seq_feature(pro_seq):
    pro_hot = np.zeros((len(pro_seq), len(pro_res_table)))
    pro_property = np.zeros((len(pro_seq), 12))
    for i in range(len(pro_seq)):
        # if 'X' in pro_seq:
        #     print(pro_seq)
        pro_hot[i,] = one_of_k_encoding(pro_seq[i], pro_res_table)
        pro_property[i,] = residue_features(pro_seq[i])
    return np.concatenate((pro_hot, pro_property), axis=1)

def target_feature(aln_file, pro_seq):
    pssm = PSSM_calculation(aln_file, pro_seq)
    other_feature = seq_feature(pro_seq)
    # print('target_feature')
    # print(pssm.shape)
    # print(other_feature.shape)

    # print(other_feature.shape)
    # return other_feature
    return np.concatenate((np.transpose(pssm, (1, 0)), other_feature), axis=1)

# target aln file save in data/dataset/aln
def target_to_feature(target_key, target_sequence, aln_dir):
    # aln_dir = 'data/' + dataset + '/aln'
    aln_file = os.path.join(aln_dir, target_key + '.aln')
    # if 'X' in target_sequence:
    #     print(target_key)
    feature = target_feature(aln_file, target_sequence)
    return feature

# pconsc4 predicted contact map save in data/dataset/pconsc4
def target_to_graph(target_key, target_sequence, contact_dir, aln_dir):
    target_size = len(target_sequence)
    contact_file = os.path.join(contact_dir, target_key + '.npy')
    # contact_dir = 'data/' + dataset + '/pconsc4'

    target_feature = target_to_feature(target_key, target_sequence, aln_dir)

    contact_map = np.load(contact_file)
    contact_map += np.matrix(np.eye(target_size))
    index_row, index_col = np.where(contact_map >= 0.5)
    target_edge_index = []
    for i, j in zip(index_row, index_col):
        target_edge_index.append([i, j])
    target_edge_index = np.array(target_edge_index)
    return target_size, target_feature, target_edge_index

def getDrugMolecularGraph(ligands):

    # create molecule graph
    print("create molecule graph ...")
    smile_graph = OrderedDict()
    for d in ligands.keys():
        lg = Chem.MolToSmiles(Chem.MolFromSmiles(ligands[d]), isomericSmiles=True)
        smile_graph[d] = smile_to_graph(lg)

    # count the number of proteins with aln and contact files
    print('effective drugs:', len(smile_graph))
    if len(smile_graph) == 0:
        raise Exception('no drug, run the script for datasets preparation.')

    return smile_graph

def getTargetMolecularGraph(proteins, dataset):

    # load contact and aln
    msa_path = 'data/' + dataset + '/aln'
    contac_path = 'data/' + dataset + '/pconsc4'

    # create protein graph
    print("create protein graph ...")
    target_graph = OrderedDict()
    for t in proteins.keys():
        g = target_to_graph(t, proteins[t], contac_path, msa_path)
        target_graph[t] = g

    # count the number of proteins with aln and contact files
    print('effective protein:', len(target_graph))
    if len(target_graph) == 0:
        raise Exception('no protein, run the script for datasets preparation.')

    return target_graph

def minMaxNormalize(Y, Y_min=None, Y_max=None):
    if Y_min is None:
        Y_min = np.min(Y)
    if Y_max is None:
        Y_max = np.max(Y)
    normalize_Y = (Y - Y_min) / (Y_max - Y_min)
    return normalize_Y

def denseAffinityRefine(adj, k):
    refine_adj = np.zeros_like(adj)
    indexs1 = np.tile(np.expand_dims(np.arange(adj.shape[0]), 0), (k, 1)).transpose()
    indexs2 = np.argpartition(adj, -k, 1)[:, -k:]
    refine_adj[indexs1, indexs2] = adj[indexs1, indexs2]
    return refine_adj

def getAffinityGraph(dataset, adj, weighted, drug_aff_k, target_aff_k):
    num_drugs = adj.shape[0]
    num_targets = adj.shape[1]
    if dataset == "davis":
        adj[adj != 0] -= 5
        adj_norm = minMaxNormalize(adj, 0)
    elif dataset == "kiba":
        adj_refine = denseAffinityRefine(adj.T, target_aff_k)
        adj_refine = denseAffinityRefine(adj_refine.T, drug_aff_k)
        adj_norm = minMaxNormalize(adj_refine, 0)
    adj_1 = adj_norm
    adj_2 = adj_norm.T

    adj = np.concatenate((
        np.concatenate((np.zeros([num_drugs, num_drugs]), adj_1), 1),
        np.concatenate((adj_2, np.zeros([num_targets, num_targets])), 1)
    ), 0)

    train_raw_ids, train_col_ids = np.where(adj != 0)
    edge_indexs = np.concatenate((
        np.expand_dims(train_raw_ids, 0),
        np.expand_dims(train_col_ids, 0)
    ), 0)
    edge_weights = adj[train_raw_ids, train_col_ids]

    node_type_features = np.concatenate((
        np.tile(np.array([1, 0]), (num_drugs, 1)),
        np.tile(np.array([0, 1]), (num_targets, 1))
    ), 0)

    adj_features = np.zeros_like(adj)
    adj_features[adj != 0] = 1

    features = np.concatenate((node_type_features, adj_features), 1)

    affinity_graph = DATA.Data(x=torch.Tensor(features), adj=torch.Tensor(adj), edge_index=torch.LongTensor(edge_indexs)) if weighted \
        else DATA.Data(x=torch.Tensor(features), adj=torch.Tensor(adj_features), edge_index=torch.LongTensor(edge_indexs))
    affinity_graph.__setitem__("edge_weight", torch.Tensor(edge_weights))
    affinity_graph.__setitem__("num_node1s", num_drugs)
    affinity_graph.__setitem__("num_node2s", num_targets)

    return affinity_graph


import numpy as np
from sacred import Ingredient

config_ingredient = Ingredient("cfg")

@config_ingredient.config
def cfg():
    # Base configuration
    model_config = {"musdb_path" : "/home/daniel/Datasets/MUSDB18", # SET MUSDB PATH HERE, AND SET CCMIXTER PATH IN CCMixter.xml
                    "estimates_path" : "/mnt/windaten/Source_Estimates", # SET THIS PATH TO WHERE YOU WANT SOURCE ESTIMATES PRODUCED BY THE TRAINED MODEL TO BE SAVED. Folder itself must exist!

                    "model_base_dir" : "checkpoints", # Base folder for model checkpoints
                    "log_dir" : "logs", # Base folder for logs files
                    "batch_size" : 16, # Batch size
                    "init_sup_sep_lr" : 1e-4, # Supervised separator learning rate
                    "epoch_it" : 2000, # Number of supervised separator steps per epoch
                    'cache_size' : 16, # Number of audio excerpts that are cached to build batches from
                    'num_workers' : 6, # Number of processes reading audio and filling up the cache
                    "duration" : 2, # Duration in seconds of the audio excerpts in the cache. Has to be at least the output length of the network!
                    'min_replacement_rate' : 16,  # roughly: how many cache entries to replace at least per batch on average. Can be fractional
                    'num_layers' : 12, # How many U-Net layers
                    'filter_size' : 15, # For Wave-U-Net: Filter size of conv in downsampling block
                    'merge_filter_size' : 5, # For Wave-U-Net: Filter size of conv in upsampling block
                    'num_initial_filters' : 24, # Number of filters for convolution in first layer of network
                    "num_frames": 16384, # DESIRED number of time frames in the output waveform per samples (could be changed when using valid padding)
                    'expected_sr': 22050,  # Downsample all audio input to this sampling rate
                    'mono_downmix': True,  # Whether to downsample the audio input
                    'output_type' : 'direct', # Type of output layer, either "direct" or "difference". Direct output: Each source is result of tanh activation and independent. DIfference: Last source output is equal to mixture input - sum(all other sources)
                    'context' : False, # Type of padding for convolutions in separator. If False, feature maps double or half in dimensions after each convolution, and convolutions are padded with zeros ("same" padding). If True, convolution is only performed on the available mixture input, thus the output is smaller than the input
                    'network' : 'unet', # Type of network architecture, either unet (our model) or unet_spectrogram (Jansson et al 2017 model)
                    'upsampling' : 'linear', # Type of technique used for upsampling the feature maps in a unet architecture, either 'linear' interpolation or 'learned' filling in of extra samples
                    'task' : 'voice', # Type of separation task. 'voice' : Separate music into voice and accompaniment. 'multi_instrument': Separate music into guitar, bass, vocals, drums and other (Sisec)
                    'augmentation' : True, # Random attenuation of source signals to improve generalisation performance (data augmentation)
                    'raw_audio_loss' : True, # Only active for unet_spectrogram network. True: L2 loss on audio. False: L1 loss on spectrogram magnitudes for training and validation and test loss
                    'worse_epochs' : 20, # Patience for early stoppping on validation set
                    }
    seed=1337
    experiment_id = np.random.randint(0,1000000)

    model_config["num_sources"] = 4 if model_config["task"] == "multi_instrument" else 2
    model_config["num_channels"] = 1 if model_config["mono_downmix"] else 2

@config_ingredient.named_config
def baseline():
    print("Training baseline model")

@config_ingredient.named_config
def baseline_diff():
    print("Training baseline model with difference output")
    model_config = {
        "output_type" : "difference"
    }

@config_ingredient.named_config
def baseline_context():
    print("Training baseline model with difference output and input context (valid convolutions)")
    model_config = {
        "output_type" : "difference",
        "context" : True
    }

@config_ingredient.named_config
def baseline_stereo():
    print("Training baseline model with difference output and input context (valid convolutions)")
    model_config = {
        "output_type" : "difference",
        "context" : True,
        "mono_downmix" : False
    }

@config_ingredient.named_config
def full():
    print("Training full singing voice separation model, with difference output and input context (valid convolutions) and stereo input/output, and learned upsampling layer")
    model_config = {
        "output_type" : "difference",
        "context" : True,
        "upsampling": "learned",
        "mono_downmix" : False
    }

@config_ingredient.named_config
def full_44KHz():
    print("Training full singing voice separation model, with difference output and input context (valid convolutions) and stereo input/output, and learned upsampling layer, and 44.1 KHz sampling rate")
    model_config = {
        "output_type" : "difference",
        "context" : True,
        "upsampling": "learned",
        "mono_downmix" : False,
        "expected_sr" : 44100
    }

@config_ingredient.named_config
def baseline_context_smallfilter_deep():
    model_config = {
        "output_type": "difference",
        "context": True,
        "num_layers" : 14,
        "duration" : 7,
        "filter_size" : 5,
        "merge_filter_size" : 1
    }

@config_ingredient.named_config
def full_multi_instrument():
    print("Training multi-instrument separation with best model")
    model_config = {
        "output_type": "difference",
        "context": True,
        "upsampling": "linear",
        "mono_downmix": False,
        "task" : "multi_instrument"
    }

@config_ingredient.named_config
def baseline_comparison():
    model_config = {
        "batch_size": 4, # Less output since model is so big. Doesn't matter since the model's output is not dependent on its output or input size (only convolutions)
        "cache_size": 4,
        "min_replacement_rate" : 4,

        "output_type": "difference",
        "context": True,
        "num_frames" : 768*127 + 1024,
        "duration" : 13,
        "expected_sr" : 8192,
        "num_initial_filters" : 34
    }

@config_ingredient.named_config
def unet_spectrogram():
    model_config = {
        "batch_size": 4, # Less output since model is so big.
        "cache_size": 4,
        "min_replacement_rate" : 4,

        "network" : "unet_spectrogram",
        "num_layers" : 6,
        "expected_sr" : 8192,
        "num_frames" : 768 * 127 + 1024, # hop_size * (time_frames_of_spectrogram_input - 1) + fft_length
        "duration" : 13,
        "num_initial_filters" : 16
    }

@config_ingredient.named_config
def unet_spectrogram_l1():
    model_config = {
        "batch_size": 4, # Less output since model is so big.
        "cache_size": 4,
        "min_replacement_rate" : 4,

        "network" : "unet_spectrogram",
        "num_layers" : 6,
        "expected_sr" : 8192,
        "num_frames" : 768 * 127 + 1024, # hop_size * (time_frames_of_spectrogram_input - 1) + fft_length
        "duration" : 13,
        "num_initial_filters" : 16,
        "loss" : "magnitudes"
    }


"""
Afterglow Core: batch photometry job plugin
"""

from datetime import datetime
from typing import List as TList, Tuple

from marshmallow.fields import Integer, List, Nested
from numpy import clip, concatenate, cos, deg2rad, isfinite, sin, zeros
from astropy.wcs import WCS
import sep

from skylib.photometry import aperture_photometry
from skylib.extraction.centroiding import centroid_sources

from ...models import (
    Job, JobResult, SourceExtractionData, PhotSettings, PhotometryData,
    sigma_to_fwhm)
from ..data_files import (
    get_data_file_data, get_exp_length, get_gain, get_image_time)

__all__ = ['PhotometryJob', 'get_source_xy', 'run_photometry_job']

def get_source_xy(source: SourceExtractionData, epoch: datetime, wcs: WCS) \
        -> Tuple[float, float]:
    """
    Calculate XY coordinates of a source in the current image, possibly taking
    proper motion into account

    :param source: source definition
    :param epoch: exposure start time
    :param wcs: WCS structure from image header

    :return: XY coordinates of the source, 1-based
    """
    if None not in (getattr(source, 'ra_hours', None),
                    getattr(source, 'dec_degs', None), wcs):
        # Prefer RA/Dec if WCS is present
        ra, dec = source.ra_hours*15, source.dec_degs
        if epoch is not None and None not in [
                getattr(source, name, None)
                for name in ('pm_sky', 'pm_pos_angle_sky', 'pm_epoch')]:
            mu = source.pm_sky*(epoch - source.pm_epoch).total_seconds()
            theta = deg2rad(source.pm_pos_angle_sky)
            cd = cos(deg2rad(dec))
            return wcs.all_world2pix(
                ((ra + mu*sin(theta)/cd) if cd else ra) % 360,
                clip(dec + mu*cos(theta), -90, 90), 1)
        return wcs.all_world2pix(ra, dec, 1)

    if None not in [getattr(source, name, None)
                    for name in ('pm_pixel', 'pm_pos_angle_pixel',
                                 'pm_epoch', 'epoch')]:
        mu = source.pm_pixel*(epoch - source.pm_epoch).total_seconds()
        theta = deg2rad(source.pm_pos_angle_pixel)
        return source.x + mu*cos(theta), source.y + mu*sin(theta)

    return source.x, source.y

def run_photometry_job(job: Job, settings: PhotSettings,
                       job_file_ids: TList[int],
                       job_sources: TList[SourceExtractionData]) \
        -> TList[PhotometryData]:
    """
    Batch photometry job body; also used during photometric calibration

    :param job: job class instance
    :param settings: photometry settings
    :param job_file_ids: data file IDs to process
    :param job_sources: list of SourceExtractionData-compatible source defs

    :return: list of photometry results
    """
    if not job_sources:
        return []

    if settings.mode == 'aperture':
        # Fixed-aperture photometry
        if settings.a is None:
            raise ValueError(
                'Missing aperture radius/semi-major axis for mode="aperture"')
        if settings.a <= 0:
            raise ValueError('Aperture radius/semi-major axis must be positive')
        phot_kw = dict(
            a=settings.a,
            b=settings.b,
            theta=settings.theta,
            a_in=settings.a_in,
            a_out=settings.a_out,
            b_out=settings.b_out,
            theta_out=settings.theta_out,
        )
    elif settings.mode == 'auto':
        # Automatic (Kron-like) photometry
        phot_kw = dict(
            k=settings.a if settings.a > 0 else 2.5,
            k_in=settings.a_in,
            k_out=settings.a_out,
        )
    else:
        raise ValueError('Photometry mode must be "aperture" or "auto"')

    # Extract file IDs from sources
    file_ids = {source.file_id for source in job_sources
                if getattr(source, 'file_id', None) is not None}
    if not file_ids:
        # Same source object for all images specified in file_ids;
        # replicate each source to all images; merge them by assigning the
        # same source ID
        if not job_file_ids and not file_ids:
            raise ValueError('Missing data file IDs')
        if job_file_ids:
            file_ids = set(job_file_ids)
        prefix = '{}_{}_'.format(
            datetime.utcnow().strftime('%Y%m%d%H%M%S'), job.id)
        sources = {
            file_id: [
                SourceExtractionData(
                    source=source, file_id=file_id,
                    id=source.id if hasattr(source, 'id') and source.id
                    else prefix + str(i + 1))
                for i, source in enumerate(job_sources)
            ] for file_id in file_ids
        }
    else:
        # Individual source object for each image; ignore job_file_ids
        if any(getattr(source, 'file_id', None) is None
               for source in job_sources):
            raise ValueError('Missing data file ID for at least one source')
        sources = {}
        for source in job_sources:
            sources.setdefault(source.file_id, []).append(source)

    # Apply custom zero point to instrumental mags if requested
    m0 = getattr(settings, 'zero_point', None) or 0

    result_data = []
    for file_no, file_id in enumerate(file_ids):
        try:
            data, hdr = get_data_file_data(job.user_id, file_id)

            if settings.gain is None:
                gain = get_gain(hdr)
            else:
                gain = settings.gain
            if gain:
                phot_kw['gain'] = gain

            epoch = get_image_time(hdr)
            texp = get_exp_length(hdr)
            flt = hdr.get('FILTER')
            scope = hdr.get('TELESCOP')

            if texp:
                phot_kw['texp'] = texp

            # noinspection PyBroadException
            try:
                wcs = WCS(hdr)
                if not wcs.has_celestial:
                    wcs = None
            except Exception:
                wcs = None

            source_table = zeros(
                len(sources[file_id]),
                [('x', float), ('y', float), ('a', float), ('b', float),
                 ('theta', float), ('flux', float), ('saturated', int),
                 ('flag', int)])
            i = 0
            while i < len(sources[file_id]):
                x, y = get_source_xy(sources[file_id][i], epoch, wcs)
                if 0 <= x < data.shape[1] and 0 <= y < data.shape[0]:
                    source_table[i]['x'], source_table[i]['y'] = x, y
                    i += 1
                else:
                    # Source outside image boundaries, skip
                    del sources[file_id][i]
                    if i:
                        if i < len(source_table) - 1:
                            source_table = concatenate(
                                [source_table[:i], source_table[i + 1:]])
                        else:
                            source_table = source_table[:i]
                    else:
                        source_table = source_table[1:]
            if not sources[file_id]:
                raise ValueError('All sources are outside image boundaries')

            r_cent = settings.centroid_radius
            if settings.mode == 'auto':
                # We need the ellipse parameters for adaptive photometry;
                # derive them from the FWHMs if available, otherwise SkyLib
                # will estimate them by isophotal analysis
                phot_kw['radius'] = r_cent
                for i, source in enumerate(sources[file_id]):
                    row = source_table[i]
                    row['a'] = getattr(source, 'fwhm_x', None) or 0
                    row['b'] = getattr(source, 'fwhm_y', None) or 0
                    row['theta'] = getattr(source, 'theta', None) or 0
                    if (not row['a'] or not row['b']) and r_cent <= 0:
                        raise ValueError(
                            'Centroiding radius must be provided for '
                            'adaptive photometry with no FWHM info')
                source_table['a'] /= sigma_to_fwhm
                source_table['b'] /= sigma_to_fwhm
                source_table['theta'] = deg2rad(source_table['theta'])

            if r_cent > 0:
                # Obtain the accurate centroids if requested
                source_table['x'], source_table['y'] = centroid_sources(
                    data, source_table['x'], source_table['y'], r_cent)

            # Photometer all sources in the current image
            source_table = aperture_photometry(data, source_table, **phot_kw)

            result_data += [
                PhotometryData(
                    source=source,
                    row=row,
                    time=epoch,
                    filter=flt,
                    telescope=scope,
                    exp_length=texp,
                    zero_point=m0,
                )
                for row, source in zip(source_table, sources[file_id])
                if row['flag'] & (0xF0 & ~sep.APER_HASMASKED) == 0 and
                isfinite([row['x'], row['y'], row['flux'], row['flux_err'],
                          row['mag'], row['mag_err']]).all()]
            job.update_progress((file_no + 1)/len(file_ids)*100)
        except Exception as e:
            job.add_error('Data file ID {}: {}'.format(file_id, e))

    return result_data

class PhotometryJobResult(JobResult):
    data: TList[PhotometryData] = List(Nested(PhotometryData), default=[])

class PhotometryJob(Job):
    type = 'photometry'
    description = 'Photometer Sources'

    result: PhotometryJobResult = Nested(PhotometryJobResult, default={})
    file_ids: TList[int] = List(Integer(), default=[])
    sources: TList[SourceExtractionData] = List(
        Nested(SourceExtractionData), default=[])
    settings: PhotSettings = Nested(PhotSettings, default={})

    def run(self):
        self.result.data = run_photometry_job(
            self, self.settings, self.file_ids, self.sources)


# Licensed under Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode)

import cv2
import numpy as np
import torch
import torchvision.transforms as transforms
from torch.backends import cudnn as cudnn
from tqdm import tqdm

from rt_gene.download_tools import download_external_landmark_models

# noinspection PyUnresolvedReferences
from rt_gene import gaze_tools as gaze_tools
from rt_gene.SFD.sfd_detector import SFDDetector
from rt_gene.ThreeDDFA.ddfa import ToTensorGjz, NormalizeGjz
from rt_gene.ThreeDDFA.inference import crop_img, predict_68pts, parse_roi_box_from_bbox, parse_roi_box_from_landmark
from rt_gene.tracker_generic import TrackedSubject

facial_landmark_transform = transforms.Compose([ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])

class LandmarkMethodBase(object):
    def __init__(self, device_id_facedetection, checkpoint_path_face=None, checkpoint_path_landmark=None, model_points_file=None):
        download_external_landmark_models()
        self.model_size_rescale = 16.0
        self.head_pitch = 0.0
        self.interpupillary_distance = 0.058
        self.eye_image_size = (60, 36)

        tqdm.write("Using device {} for face detection.".format(device_id_facedetection))

        self.device = device_id_facedetection
        self.face_net = SFDDetector(device=device_id_facedetection, path_to_detector=checkpoint_path_face)
        self.facial_landmark_nn = self.load_face_landmark_model(checkpoint_path_landmark)

        self.model_points = self.get_full_model_points(model_points_file)

    def load_face_landmark_model(self, checkpoint_fp=None):
        import rt_gene.ThreeDDFA.mobilenet_v1 as mobilenet_v1
        if checkpoint_fp is None:
            import rospkg
            checkpoint_fp = rospkg.RosPack().get_path('rt_gene') + '/model_nets/phase1_wpdc_vdc.pth.tar'
        arch = 'mobilenet_1'

        checkpoint = torch.load(checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
        model = getattr(mobilenet_v1, arch)(num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

        model_dict = model.state_dict()
        # because the model is trained by multiple gpus, prefix module should be removed
        for k in checkpoint.keys():
            model_dict[k.replace('module.', '')] = checkpoint[k]
        model.load_state_dict(model_dict)
        cudnn.benchmark = True
        model = model.to(self.device)
        model.eval()
        return model

    def get_full_model_points(self, model_points_file=None):
        """Get all 68 3D model points from file"""
        raw_value = []
        if model_points_file is None:
            import rospkg
            model_points_file = rospkg.RosPack().get_path('rt_gene') + '/model_nets/face_model_68.txt'

        with open(model_points_file) as f:
            for line in f:
                raw_value.append(line)
        model_points = np.array(raw_value, dtype=np.float32)
        model_points = np.reshape(model_points, (3, -1)).T

        # index the expansion of the model based.
        model_points = model_points * (self.interpupillary_distance * self.model_size_rescale)

        return model_points

    def get_face_bb(self, image):
        faceboxes = []
        fraction = 4.0
        image = cv2.resize(image, (0, 0), fx=1.0 / fraction, fy=1.0 / fraction)
        detections = self.face_net.detect_from_image(image)

        for result in detections:
            # scale back up to image size
            box = result[:4]
            confidence = result[4]

            if gaze_tools.box_in_image(box, image) and confidence > 0.6:
                box = [x * fraction for x in box]  # scale back up
                diff_height_width = (box[3] - box[1]) - (box[2] - box[0])
                offset_y = int(abs(diff_height_width / 2))
                box_moved = gaze_tools.move_box(box, [0, offset_y])

                # Make box square.
                facebox = gaze_tools.get_square_box(box_moved)
                faceboxes.append(facebox)

        return faceboxes

    @staticmethod
    def visualize_headpose_result(face_image, est_headpose):
        """Here, we take the original eye eye_image and overlay the estimated headpose."""
        output_image = np.copy(face_image)

        center_x = output_image.shape[1] / 2
        center_y = output_image.shape[0] / 2

        endpoint_x, endpoint_y = gaze_tools.get_endpoint(est_headpose[1], est_headpose[0], center_x, center_y, 100)

        cv2.line(output_image, (int(center_x), int(center_y)), (int(endpoint_x), int(endpoint_y)), (0, 0, 255), 3)
        return output_image

    def ddfa_forward_pass(self, color_img, roi_box_list):
        img_step = [crop_img(color_img, roi_box) for roi_box in roi_box_list]
        img_step = [cv2.resize(img, dsize=(120, 120), interpolation=cv2.INTER_LINEAR) for img in img_step]
        _input = torch.cat([facial_landmark_transform(img).unsqueeze(0) for img in img_step], 0)
        with torch.no_grad():
            _input = _input.to(self.device)
            param = self.facial_landmark_nn(_input).cpu().numpy().astype(np.float32)

        return [predict_68pts(p.flatten(), roi_box) for p, roi_box in zip(param, roi_box_list)]

    def get_subjects_from_faceboxes(self, color_img, faceboxes):
        face_images = [gaze_tools.crop_face_from_image(color_img, b) for b in faceboxes]
        subjects = []
        roi_box_list = [parse_roi_box_from_bbox(facebox) for facebox in faceboxes]
        initial_pts68_list = self.ddfa_forward_pass(color_img, roi_box_list)
        roi_box_refined_list = [parse_roi_box_from_landmark(initial_pts68) for initial_pts68 in initial_pts68_list]
        pts68_list = self.ddfa_forward_pass(color_img, roi_box_refined_list)

        for pts68, face_image, facebox in zip(pts68_list, face_images, faceboxes):
            np_landmarks = np.array((pts68[0], pts68[1])).T
            subjects.append(TrackedSubject(np.array(facebox), face_image, np_landmarks))
        return subjects


"""
Code for running a 1D-Convolutional neural network on SMILES strings
Derek van Tilborg, Eindhoven University of Technology, March 2022

inspired by:

Kimber, T. B., Gagnebin, M. & Volkamer, A. Maxsmi: Maximizing molecular property prediction performance with confidence
estimation using SMILES augmentation and deep learning. Artificial Intelligence in the Life Sciences 1, 100014 (2021).

SMILES encoding settings are adapted from:

Moret, M., Grisoni, F., Katzberger, P. & Schneider, G.
Perplexity-based molecule ranking and bias estimation of chemical language models.
ChemRxiv (2021) doi:10.26434/chemrxiv-2021-zv6f1-v2.
"""

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten, Convolution1D, Dropout, BatchNormalization, MaxPooling1D
import tensorflow.keras.optimizers
from tensorflow.keras.callbacks import ReduceLROnPlateau, EarlyStopping
from tensorflow.random import set_seed
import numpy as np

from MoleculeACE.benchmark import utils
from MoleculeACE.benchmark.utils.const import RANDOM_SEED
from MoleculeACE.benchmark.utils.const import CONFIG_PATH_GENERAL, CONFIG_PATH_SMILES, CONFIG_PATH_CNN
from MoleculeACE.benchmark.data_processing.preprocessing.one_hot_encoding import onehot_to_smiles
from MoleculeACE.benchmark.data_processing.preprocessing.data_prep import split_smiles

cnn_settings = utils.get_config(CONFIG_PATH_CNN)
smiles_encoding = utils.get_config(CONFIG_PATH_SMILES)
general_settings = utils.get_config(CONFIG_PATH_GENERAL)

class SeqModel:
    """ Define and build the 1D CNN """
    def __init__(self, smiles_length=202, n_chars=35, dropout=0.5, lr=0.0001, batch_size=32):
        self.model = None
        self.smiles_length = smiles_length
        self.n_chars = n_chars
        self.dropout = dropout
        self.lr = lr
        self.batch_size = batch_size

        self.build_model()

    def build_model(self):
        self.model = Sequential()
        self.model.add(Convolution1D(32, 1, activation="relu", input_shape=(self.smiles_length, self.n_chars)))
        self.model.add(MaxPooling1D(2))

        self.model.add(Flatten())
        self.model.add(BatchNormalization())
        self.model.add(Dropout(self.dropout))
        self.model.add(Dense(128, activation="relu"))
        self.model.add(Dense(1))

        optimizer = tensorflow.keras.optimizers.Adam(learning_rate=self.lr)
        self.model.compile(loss='mean_squared_error', optimizer=optimizer)

def smiles_cnn(x_train, y_train, x_val=None, y_val=None, save_path='Results/'):
    """ Train a 1D CNN on one-hot encoded smiles data

    Args:
        x_train: (array) one-hot encoded smiles with shape (n_samples, max_smiles_len, n_characters)
        y_train: (lst) List of bioactivity values
        y_val: (array) one-hot encoded validation smiles with shape (n_samples, max_smiles_len, n_char) (default=None)
        x_val: (lst) List of validation bioactivity values (default=None)
        epochs: (int) number of epochs the model will be trained
        save_path: (str) path where best models are stored

    Returns: MoleculeACE.CNN.model.SeqModel

    """

    # Get the hyperparameters from the cnn config file
    dropout = cnn_settings['dropout']
    lr = cnn_settings['lr']
    batch_size = cnn_settings['batch_size']
    patience_stopping = cnn_settings['early_stopping_patience']
    monitor = cnn_settings['monitor']
    patience_lr = cnn_settings['patience_lr']
    lr_factor = cnn_settings['lr_factor']
    min_lr = cnn_settings['min_lr']
    n_workers = general_settings['num_workers']
    val_split = cnn_settings['val_split']

    # Set random seed
    set_seed(RANDOM_SEED)

    # convert the list of train labels to an array
    y_train = np.array(y_train)
    if y_val is not None:
        y_val = np.array(y_val)

    # find smiles array dimensions
    smiles_length = x_train.shape[1]
    n_chars = x_train.shape[2]

    # Callbacks
    early_stopping = EarlyStopping(monitor='val_loss', mode='min', verbose=1,
                                   patience=patience_stopping)
    checkpointer = utils.create_model_checkpoint(f"{save_path}/cnn_")
    lr_reduction = ReduceLROnPlateau(monitor=monitor,
                                     patience=patience_lr,
                                     factor=lr_factor,
                                     min_lr=min_lr)

    # Initiate and compile the model
    model = SeqModel(smiles_length=smiles_length, n_chars=n_chars, dropout=dropout, lr=lr, batch_size=batch_size).model

    # Train the model
    if x_val is None and y_val is None:
        print(f"Creating a validation split from the train data ({val_split})")
        train_smiles = []
        for array in x_train:
            smi = onehot_to_smiles(array, smiles_encoding['indices_token'],
                                   smiles_encoding['start_char'],
                                   smiles_encoding['end_char'])
            train_smiles.append(smi)

        # Split based on Murco scaffold clustering
        train_idx, val_idx = split_smiles(train_smiles, test_split=val_split, clustering_cutoff=0.4)

        # Split the train data into train and validation
        x_val = x_train[val_idx]
        x_train = x_train[train_idx]
        y_val = y_train[val_idx]
        y_train = y_train[train_idx]

    history = model.fit(x=x_train, y=y_train, validation_data=(x_val, y_val), use_multiprocessing=True,
                              epochs=cnn_settings['epochs'],
                              callbacks=[checkpointer, lr_reduction, early_stopping],
                              workers=n_workers, verbose=1)

    # plot the train/val loss
    utils.plot_history(history, 'MSE')

    return model


import numpy as np
import os,xml.etree.ElementTree as ET
import cv2
import matplotlib.pyplot as plt
from glob import glob
from concurrent.futures import ThreadPoolExecutor
import tensorflow as tf
from PIL import Image,ImageDraw
from math import *
anchor_scale = 16
#
IOU_NEGATIVE = 0.3
IOU_POSITIVE = 0.7
IOU_SELECT = 0.7

RPN_POSITIVE_NUM = 150
RPN_TOTAL_NUM = 300

# bgr  can find from  here https://github.com/fchollet/deep-learning-models/blob/master/imagenet_utils.py
IMAGE_MEAN = [123.68, 116.779, 103.939]

DEBUG = True
def rotate(img_path,degree=90):

    img=cv2.imread(img_path)
    height, width = img.shape[:2]

    # 旋转后的尺寸
    heightNew = int(width * fabs(sin(radians(degree))) + height * fabs(cos(radians(degree))))  # 这个公式参考之前内容
    widthNew = int(height * fabs(sin(radians(degree))) + width * fabs(cos(radians(degree))))

    matRotation = cv2.getRotationMatrix2D((width / 2, height / 2), degree, 1)

    matRotation[0, 2] += (widthNew - width) / 2  # 因为旋转之后,坐标系原点是新图像的左上角,所以需要根据原图做转化
    matRotation[1, 2] += (heightNew - height) / 2

    imgRotation = cv2.warpAffine(img, matRotation, (widthNew, heightNew), borderValue=(255, 255, 255))

    return Image.fromarray(imgRotation)

def drawRect(gtboxes,img):
    draw=ImageDraw.Draw(img)
    for i in gtboxes:
        draw.rectangle((i[0],i[1],i[2],i[3]),outline='red')
    img.show()

def readxml(path):
    gtboxes=[]
    tree=ET.parse(path)
    root=tree.getroot()
    image_name = tree.find('filename').text
    for obj in root.iter('object'):
        try:
            bnd=obj.find('bndbox')
            xmin=int(bnd.find('xmin').text)
            ymin=int(bnd.find('ymin').text)
            xmax=int(bnd.find('xmax').text)
            ymax=int(bnd.find('ymax').text)
            gtboxes.append((xmin,ymin,xmax,ymax))
        except Exception as e:
            raise Exception(e+" when {} loading xml file, there are errors".format(path))
    return np.array(gtboxes),image_name

def gen_anchor(featuresize, scale):
    """
    gen base anchor from feature map [HXW][10][4]
    reshape  [HXW][10][4] to [HXWX10][4]
    生成的锚框是相对于原图的，即原图中每16像素就有10个锚框
    """

    heights = [11, 16, 23, 33, 48, 68, 97, 139, 198, 283]
    widths = [16, 16, 16, 16, 16, 16, 16, 16, 16, 16]

    # gen k=9 anchor size (h,w)
    heights = np.array(heights).reshape(len(heights), 1)
    widths = np.array(widths).reshape(len(widths), 1)

    # 锚框大小为16像素
    base_anchor = np.array([0, 0, 15, 15])
    # center x,y
    xt = (base_anchor[0] + base_anchor[2]) * 0.5
    yt = (base_anchor[1] + base_anchor[3]) * 0.5

    # x1 y1 x2 y2
    x1 = xt - widths * 0.5
    y1 = yt - heights * 0.5
    x2 = xt + widths * 0.5
    y2 = yt + heights * 0.5

    # 一组十个锚框
    base_anchor = np.hstack((x1, y1, x2, y2))

    h, w = featuresize
    shift_x = np.arange(0, w) * scale
    shift_y = np.arange(0, h) * scale
    # apply shift
    anchor = []
    for i in shift_y:
        for j in shift_x:
            anchor.append(base_anchor + [j, i, j, i])
    return np.array(anchor).reshape((-1, 4))

def cal_iou(box1, box1_area, boxes2, boxes2_area):
    """
    box1 [x1,y1,x2,y2]
    boxes2 [Msample,x1,y1,x2,y2]
    """
    x1 = np.maximum(box1[0], boxes2[:, 0])
    x2 = np.minimum(box1[2], boxes2[:, 2])
    y1 = np.maximum(box1[1], boxes2[:, 1])
    y2 = np.minimum(box1[3], boxes2[:, 3])

    intersection = np.maximum(x2 - x1, 0) * np.maximum(y2 - y1, 0)
    iou = intersection / (box1_area + boxes2_area[:] - intersection[:])
    return iou

def cal_overlaps(boxes1, boxes2):
    """
    boxes1 [Nsample,x1,y1,x2,y2]  anchor
    boxes2 [Msample,x1,y1,x2,y2]  grouth-box

    """
    area1 = (boxes1[:, 0] - boxes1[:, 2]) * (boxes1[:, 1] - boxes1[:, 3])  # (Nsample, 1)
    area2 = (boxes2[:, 0] - boxes2[:, 2]) * (boxes2[:, 1] - boxes2[:, 3])  # (Msample, 1)

    overlaps = np.zeros((boxes1.shape[0], boxes2.shape[0]))  # (Nsample, Msample)

    # calculate the intersection of  boxes1(anchor) and boxes2(GT box)
    for i in range(boxes1.shape[0]):
        overlaps[i][:] = cal_iou(boxes1[i], area1[i], boxes2, area2)

    return overlaps

def bbox_transfrom(anchors, gtboxes):
    """
    anchors: (Nsample, 4)
    gtboxes: (Nsample, 4)
     compute relative predicted vertical coordinates Vc ,Vh
        with respect to the bounding box location of an anchor
    """
    Cy = (gtboxes[:, 1] + gtboxes[:, 3]) * 0.5  # (Nsample, )
    Cya = (anchors[:, 1] + anchors[:, 3]) * 0.5  # (Nsample, )
    h = gtboxes[:, 3] - gtboxes[:, 1] + 1.0  # (Nsample, )
    ha = anchors[:, 3] - anchors[:, 1] + 1.0  # (Nsample, )

    Vc = (Cy - Cya) / ha  # (Nsample, )
    Vh = np.log(h / ha)  # (Nsample, )

    ret = np.vstack((Vc, Vh))

    return ret.transpose()  # (Nsample, 2)

def bbox_transfor_inv(anchor, regr):
    """
    anchor: (NSample, 4)
    regr: (NSample, 2)
    根据锚框和偏移量反向得到GTBox
    """

    Cya = (anchor[:, 1] + anchor[:, 3]) * 0.5  # 锚框y中心点
    ha = anchor[:, 3] - anchor[:, 1] + 1

    Vcx = regr[..., 0]  # y中心点偏移
    Vhx = regr[..., 1]  # 高度偏移

    Cyx = Vcx * ha + Cya  # GTBox y中心点
    hx = np.exp(Vhx) * ha  # GTBox 高
    xt = (anchor[:, 0] + anchor[:, 2]) * 0.5  # 锚框x中心点

    x1 = xt - 16 * 0.5
    y1 = Cyx - hx * 0.5
    x2 = xt + 16 * 0.5
    y2 = Cyx + hx * 0.5
    bbox = np.vstack((x1, y1, x2, y2)).transpose()

    return bbox

def clip_box(bbox, im_shape):
    # x1 >= 0
    bbox[:, 0] = np.maximum(np.minimum(bbox[:, 0], im_shape[1] - 1), 0)
    # y1 >= 0
    bbox[:, 1] = np.maximum(np.minimum(bbox[:, 1], im_shape[0] - 1), 0)
    # x2 < im_shape[1]
    bbox[:, 2] = np.maximum(np.minimum(bbox[:, 2], im_shape[1] - 1), 0)
    # y2 < im_shape[0]
    bbox[:, 3] = np.maximum(np.minimum(bbox[:, 3], im_shape[0] - 1), 0)

    return bbox

def filter_bbox(bbox, minsize):
    ws = bbox[:, 2] - bbox[:, 0] + 1
    hs = bbox[:, 3] - bbox[:, 1] + 1
    keep = np.where((ws >= minsize) & (hs >= minsize))[0]
    return keep

def cal_rpn(imgsize, featuresize, scale, gtboxes):
    """
    gtboxes: (Msample, 4)
    """
    imgh, imgw = imgsize

    # gen base anchor
    base_anchor = gen_anchor(featuresize, scale)  # (Nsample, 4)

    # calculate iou
    overlaps = cal_overlaps(base_anchor, gtboxes)  # (Nsample, Msample)

    # init labels -1 don't care  0 is negative  1 is positive
    labels = np.empty(base_anchor.shape[0])
    labels.fill(-1)  # (Nsample,)

    # for each GT box corresponds to an anchor which has highest IOU
    gt_argmax_overlaps = overlaps.argmax(axis=0)  # (Msample, )

    # the anchor with the highest IOU overlap with a GT box
    anchor_argmax_overlaps = overlaps.argmax(axis=1)  # (Nsample, )
    anchor_max_overlaps = overlaps[range(overlaps.shape[0]), anchor_argmax_overlaps]  # (Nsample, )
    print('anchor max overlaps')
    for i in anchor_max_overlaps:
        if i>IOU_POSITIVE:
            print(i)

    # IOU > IOU_POSITIVE
    labels[anchor_max_overlaps > IOU_POSITIVE] = 1
    # IOU <IOU_NEGATIVE
    labels[anchor_max_overlaps < IOU_NEGATIVE] = 0
    # ensure that every GT box has at least one positive RPN region
    labels[gt_argmax_overlaps] = 1
    # print('内部的label',labels)
    # for i in labels:
    #     print(i)

    # only keep anchors inside the image
    outside_anchor = np.where(
        (base_anchor[:, 0] < 0) |
        (base_anchor[:, 1] < 0) |
        (base_anchor[:, 2] >= imgw) |
        (base_anchor[:, 3] >= imgh)
    )[0]
    labels[outside_anchor] = -1

    # 剔除掉多余的正负样例
    # subsample positive labels ,if greater than RPN_POSITIVE_NUM(default 128)
    fg_index = np.where(labels == 1)[0]
    if (len(fg_index) > RPN_POSITIVE_NUM):
        labels[np.random.choice(fg_index, len(fg_index) - RPN_POSITIVE_NUM, replace=False)] = -1

    # subsample negative labels
    bg_index = np.where(labels == 0)[0]
    num_bg = RPN_TOTAL_NUM - np.sum(labels == 1)
    if (len(bg_index) > num_bg):
        # print('bgindex:',len(bg_index),'num_bg',num_bg)
        labels[np.random.choice(bg_index, len(bg_index) - num_bg, replace=False)] = -1

    # calculate bbox targets
    # debug here
    bbox_targets = bbox_transfrom(base_anchor, gtboxes[anchor_argmax_overlaps, :])
    # bbox_targets=[]

    return [labels, bbox_targets], base_anchor

def get_session(gpu_fraction=0.6):
    '''''Assume that you have 6GB of GPU memory and want to allocate ~2GB'''

    num_threads = os.environ.get('OMP_NUM_THREADS')
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction)

    if num_threads:
        return tf.Session(config=tf.ConfigProto(
            gpu_options=gpu_options, intra_op_parallelism_threads=num_threads, allow_soft_placement=True))
    else:
        return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, allow_soft_placement=True))

class random_uniform_num():
    """
    uniform random
    """

    def __init__(self, total, start=0):
        self.total = total
        self.range = [i for i in range(total)]
        np.random.shuffle(self.range)
        self.index = start

    def get(self, batch_size):
        ret = []
        if self.index + batch_size > self.total:
            piece1 = self.range[self.index:]
            np.random.shuffle(self.range)
            self.index = (self.index + batch_size) - self.total
            piece2 = self.range[0:self.index]
            ret.extend(piece1)
            ret.extend(piece2)
        else:
            ret = self.range[self.index:self.index + batch_size]
            self.index = self.index + batch_size
        return ret

def nms(dets, thresh):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]  # Sort from high to low

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]
    return keep

def gen_sample(xmlpath, imgpath, batchsize=1):
    """
    由于图像大小不定，批处理大小只能为1
    """

    # list xml
    xmlfiles = glob(xmlpath + '/*.xml')
    rd = random_uniform_num(len(xmlfiles))
    xmlfiles = np.array(xmlfiles)

    while True:
        shuf = xmlfiles[rd.get(1)]
        gtbox, imgfile = readxml(shuf[0])
        img = cv2.imread(imgpath + "\\" + imgfile)
        h, w, c = img.shape

        # clip image
        if np.random.randint(0, 100) > 50:
            img = img[:, ::-1, :]
            newx1 = w - gtbox[:, 2] - 1
            newx2 = w - gtbox[:, 0] - 1
            gtbox[:, 0] = newx1
            gtbox[:, 2] = newx2

        [cls, regr], _ = cal_rpn((h, w), (int(h / 16), int(w / 16)), 16, gtbox)
        # zero-center by mean pixel
        m_img = img - IMAGE_MEAN
        m_img = np.expand_dims(m_img, axis=0)

        regr = np.hstack([cls.reshape(cls.shape[0], 1), regr])

        #
        cls = np.expand_dims(cls, axis=0)
        cls = np.expand_dims(cls, axis=1)
        # regr = np.expand_dims(regr,axis=1)
        regr = np.expand_dims(regr, axis=0)

        yield m_img, {'rpn_class_reshape': cls, 'rpn_regress_reshape': regr}

def rpn_test():
    xmlpath = 'G:\data\VOCdevkit\VOC2007\Annotations\img_4375.xml'
    imgpath = 'G:\data\VOCdevkit\VOC2007\JPEGImages\img_4375.jpg'
    gtbox, _ = readxml(xmlpath)
    img = cv2.imread(imgpath)
    h, w, c = img.shape
    [cls, regr], base_anchor = cal_rpn((h, w), (int(h / 16), int(w / 16)), 16, gtbox)
    print(cls.shape)
    print(regr.shape)

    regr = np.expand_dims(regr, axis=0)
    inv_anchor = bbox_transfor_inv(base_anchor, regr)
    anchors = inv_anchor[cls == 1]
    anchors = anchors.astype(int)
    for i in anchors:
        cv2.rectangle(img, (i[0], i[1]), (i[2], i[3]), (255, 0, 0), 3)
    plt.imshow(img)

def threshold(coords, min_, max_):
    return np.maximum(np.minimum(coords, max_), min_)

def clip_boxes(boxes, im_shape):
    """
    Clip boxes to image boundaries.
    """
    boxes[:, 0::2]=threshold(boxes[:, 0::2], 0, im_shape[1]-1)
    boxes[:, 1::2]=threshold(boxes[:, 1::2], 0, im_shape[0]-1)
    return boxes

class Graph:
    def __init__(self, graph):
        self.graph=graph

    def sub_graphs_connected(self):
        sub_graphs=[]
        for index in range(self.graph.shape[0]):
            if not self.graph[:, index].any() and self.graph[index, :].any():
                v=index
                sub_graphs.append([v])
                while self.graph[v, :].any():
                    v=np.where(self.graph[v, :])[0][0]
                    sub_graphs[-1].append(v)
        return sub_graphs
class TextLineCfg:
    SCALE=600
    MAX_SCALE=1200
    TEXT_PROPOSALS_WIDTH=16
    MIN_NUM_PROPOSALS = 2
    MIN_RATIO=0.5
    LINE_MIN_SCORE=0.9
    MAX_HORIZONTAL_GAP=60
    TEXT_PROPOSALS_MIN_SCORE=0.7
    TEXT_PROPOSALS_NMS_THRESH=0.3
    MIN_V_OVERLAPS=0.6
    MIN_SIZE_SIM=0.6

class TextProposalGraphBuilder:
    """
        Build Text proposals into a graph.
    """
    def get_successions(self, index):
            box=self.text_proposals[index]
            results=[]
            for left in range(int(box[0])+1, min(int(box[0])+TextLineCfg.MAX_HORIZONTAL_GAP+1, self.im_size[1])):
                adj_box_indices=self.boxes_table[left]
                for adj_box_index in adj_box_indices:
                    if self.meet_v_iou(adj_box_index, index):
                        results.append(adj_box_index)
                if len(results)!=0:
                    return results
            return results

    def get_precursors(self, index):
        box=self.text_proposals[index]
        results=[]
        for left in range(int(box[0])-1, max(int(box[0]-TextLineCfg.MAX_HORIZONTAL_GAP), 0)-1, -1):
            adj_box_indices=self.boxes_table[left]
            for adj_box_index in adj_box_indices:
                if self.meet_v_iou(adj_box_index, index):
                    results.append(adj_box_index)
            if len(results)!=0:
                return results
        return results

    def is_succession_node(self, index, succession_index):
        precursors=self.get_precursors(succession_index)
        if self.scores[index]>=np.max(self.scores[precursors]):
            return True
        return False

    def meet_v_iou(self, index1, index2):
        def overlaps_v(index1, index2):
            h1=self.heights[index1]
            h2=self.heights[index2]
            y0=max(self.text_proposals[index2][1], self.text_proposals[index1][1])
            y1=min(self.text_proposals[index2][3], self.text_proposals[index1][3])
            return max(0, y1-y0+1)/min(h1, h2)

        def size_similarity(index1, index2):
            h1=self.heights[index1]
            h2=self.heights[index2]
            return min(h1, h2)/max(h1, h2)

        return overlaps_v(index1, index2)>=TextLineCfg.MIN_V_OVERLAPS and \
               size_similarity(index1, index2)>=TextLineCfg.MIN_SIZE_SIM

    def build_graph(self, text_proposals, scores, im_size):
        self.text_proposals=text_proposals
        self.scores=scores
        self.im_size=im_size
        self.heights=text_proposals[:, 3]-text_proposals[:, 1]+1

        boxes_table=[[] for _ in range(self.im_size[1])]
        for index, box in enumerate(text_proposals):
            boxes_table[int(box[0])].append(index)
        self.boxes_table=boxes_table

        graph=np.zeros((text_proposals.shape[0], text_proposals.shape[0]), np.bool)

        for index, box in enumerate(text_proposals):
            successions=self.get_successions(index)
            if len(successions)==0:
                continue
            succession_index=successions[np.argmax(scores[successions])]
            if self.is_succession_node(index, succession_index):
                # NOTE: a box can have multiple successions(precursors) if multiple successions(precursors)
                # have equal scores.
                graph[index, succession_index]=True
        return Graph(graph)
class TextProposalConnector:
    def __init__(self):
        self.graph_builder=TextProposalGraphBuilder()

    def group_text_proposals(self, text_proposals, scores, im_size):
        graph=self.graph_builder.build_graph(text_proposals, scores, im_size)
        return graph.sub_graphs_connected()

    def fit_y(self, X, Y, x1, x2):
        len(X)!=0
        # if X only include one point, the function will get line y=Y[0]
        if np.sum(X==X[0])==len(X):
            return Y[0], Y[0]
        p=np.poly1d(np.polyfit(X, Y, 1))
        return p(x1), p(x2)

    def get_text_lines(self, text_proposals, scores, im_size):
        # tp=text proposal
        tp_groups=self.group_text_proposals(text_proposals, scores, im_size)
        text_lines=np.zeros((len(tp_groups), 5), np.float32)

        for index, tp_indices in enumerate(tp_groups):
            text_line_boxes=text_proposals[list(tp_indices)]

            x0=np.min(text_line_boxes[:, 0])
            x1=np.max(text_line_boxes[:, 2])

            offset=(text_line_boxes[0, 2]-text_line_boxes[0, 0])*0.5

            lt_y, rt_y=self.fit_y(text_line_boxes[:, 0], text_line_boxes[:, 1], x0+offset, x1-offset)
            lb_y, rb_y=self.fit_y(text_line_boxes[:, 0], text_line_boxes[:, 3], x0+offset, x1-offset)

            # the score of a text line is the average score of the scores
            # of all text proposals contained in the text line
            score=scores[list(tp_indices)].sum()/float(len(tp_indices))

            text_lines[index, 0]=x0
            text_lines[index, 1]=min(lt_y, rt_y)
            text_lines[index, 2]=x1
            text_lines[index, 3]=max(lb_y, rb_y)
            text_lines[index, 4]=score

        text_lines=clip_boxes(text_lines, im_size)

        text_recs = np.zeros((len(text_lines), 9), np.float)
        index = 0
        for line in text_lines:
            xmin,ymin,xmax,ymax=line[0],line[1],line[2],line[3]
            text_recs[index, 0] = xmin
            text_recs[index, 1] = ymin
            text_recs[index, 2] = xmax
            text_recs[index, 3] = ymin
            text_recs[index, 4] = xmin
            text_recs[index, 5] = ymax
            text_recs[index, 6] = xmax
            text_recs[index, 7] = ymax
            text_recs[index, 8] = line[4]
            index = index + 1

        return text_recs

class TextProposalConnectorOriented:
    """
        Connect text proposals into text lines
    """

    def __init__(self):
        self.graph_builder = TextProposalGraphBuilder()

    def group_text_proposals(self, text_proposals, scores, im_size):
        graph = self.graph_builder.build_graph(text_proposals, scores, im_size)
        return graph.sub_graphs_connected()

    def fit_y(self, X, Y, x1, x2):
        len(X) != 0
        # if X only include one point, the function will get line y=Y[0]
        if np.sum(X == X[0]) == len(X):
            return Y[0], Y[0]
        p = np.poly1d(np.polyfit(X, Y, 1))
        return p(x1), p(x2)

    def get_text_lines(self, text_proposals, scores, im_size):
        """
        text_proposals:boxes

        """
        # tp=text proposal
        tp_groups = self.group_text_proposals(text_proposals, scores, im_size)  # 首先还是建图，获取到文本行由哪几个小框构成

        text_lines = np.zeros((len(tp_groups), 8), np.float32)

        for index, tp_indices in enumerate(tp_groups):
            text_line_boxes = text_proposals[list(tp_indices)]  # 每个文本行的全部小框
            X = (text_line_boxes[:, 0] + text_line_boxes[:, 2]) / 2  # 求每一个小框的中心x，y坐标
            Y = (text_line_boxes[:, 1] + text_line_boxes[:, 3]) / 2

            z1 = np.polyfit(X, Y, 1)  # 多项式拟合，根据之前求的中心店拟合一条直线（最小二乘）

            x0 = np.min(text_line_boxes[:, 0])  # 文本行x坐标最小值
            x1 = np.max(text_line_boxes[:, 2])  # 文本行x坐标最大值

            offset = (text_line_boxes[0, 2] - text_line_boxes[0, 0]) * 0.5  # 小框宽度的一半

            # 以全部小框的左上角这个点去拟合一条直线，然后计算一下文本行x坐标的极左极右对应的y坐标
            lt_y, rt_y = self.fit_y(text_line_boxes[:, 0], text_line_boxes[:, 1], x0 + offset, x1 - offset)
            # 以全部小框的左下角这个点去拟合一条直线，然后计算一下文本行x坐标的极左极右对应的y坐标
            lb_y, rb_y = self.fit_y(text_line_boxes[:, 0], text_line_boxes[:, 3], x0 + offset, x1 - offset)

            score = scores[list(tp_indices)].sum() / float(len(tp_indices))  # 求全部小框得分的均值作为文本行的均值

            text_lines[index, 0] = x0
            text_lines[index, 1] = min(lt_y, rt_y)  # 文本行上端 线段 的y坐标的小值
            text_lines[index, 2] = x1
            text_lines[index, 3] = max(lb_y, rb_y)  # 文本行下端 线段 的y坐标的大值
            text_lines[index, 4] = score  # 文本行得分
            text_lines[index, 5] = z1[0]  # 根据中心点拟合的直线的k，b
            text_lines[index, 6] = z1[1]
            height = np.mean((text_line_boxes[:, 3] - text_line_boxes[:, 1]))  # 小框平均高度
            text_lines[index, 7] = height + 2.5

        text_recs = np.zeros((len(text_lines), 9), np.float)
        index = 0
        for line in text_lines:
            b1 = line[6] - line[7] / 2  # 根据高度和文本行中心线，求取文本行上下两条线的b值
            b2 = line[6] + line[7] / 2
            x1 = line[0]
            y1 = line[5] * line[0] + b1  # 左上
            x2 = line[2]
            y2 = line[5] * line[2] + b1  # 右上
            x3 = line[0]
            y3 = line[5] * line[0] + b2  # 左下
            x4 = line[2]
            y4 = line[5] * line[2] + b2  # 右下
            disX = x2 - x1
            disY = y2 - y1
            width = np.sqrt(disX * disX + disY * disY)  # 文本行宽度

            fTmp0 = y3 - y1  # 文本行高度
            fTmp1 = fTmp0 * disY / width
            x = np.fabs(fTmp1 * disX / width)  # 做补偿
            y = np.fabs(fTmp1 * disY / width)
            if line[5] < 0:
                x1 -= x
                y1 += y
                x4 += x
                y4 -= y
            else:
                x2 += x
                y2 += y
                x3 -= x
                y3 -= y
            text_recs[index, 0] = x1
            text_recs[index, 1] = y1
            text_recs[index, 2] = x2
            text_recs[index, 3] = y2
            text_recs[index, 4] = x3
            text_recs[index, 5] = y3
            text_recs[index, 6] = x4
            text_recs[index, 7] = y4
            text_recs[index, 8] = line[4]
            index = index + 1

        return text_recs

if __name__=='__main__':
    xmlpath = 'D:\py_projects\data_new\data_new\data\\annotation\\img_calligraphy_00003_bg.xml'
    imgpath = 'D:\py_projects\data_new\data_new\data\\train_img\\img_calligraphy_00003_bg.jpg'
    gtboxes, _ = readxml(xmlpath)
    img = rotate(imgpath,90)
    print(gtboxes)
    ow,oh=img.size
    newbox = []
    for i in gtboxes:
        newbox.append([i[1], ow - i[2], i[3], ow - i[0]])
    gtboxes = np.array(newbox)
    image_shape = (512,256)
    x_scale = image_shape[0] / ow
    y_scale = image_shape[1] / oh
    newbox = []
    for i in range(len(gtboxes)):
        newbox.append(
            [gtboxes[i][0] * x_scale, gtboxes[i][1] * y_scale, gtboxes[i][2] * x_scale,
             gtboxes[i][3] * y_scale]
        )
    img = img.resize(image_shape, Image.ANTIALIAS)
    gtboxes = np.array(newbox)
    print(gtboxes)
    [cls,regr],base_anchor=cal_rpn((256,512),(256//16,512//16),16,gtboxes)
    print(cls)
    regr=np.expand_dims(regr,axis=0)
    anchor=bbox_transfor_inv(base_anchor,regr)
    anchor=anchor[cls==1]
    anchor=anchor.astype(int)
    print(anchor)
    drawRect(anchor,img)