"""
deal with deep learning models
"""
from tqdm import tqdm

from keras_bert import load_trained_model_from_checkpoint, Tokenizer
from keras.layers import *
from keras.models import Model
from keras.optimizers import Adam
from keras_attention import SelfAttention
import keras.backend as K

from optparse import OptionParser

from ie_utils import *

class piecewise_maxpool_layer(Layer):
    def __init__(self, filter_num, seq_len, **kwargs):
        self.filter_num = filter_num
        self.seq_len = seq_len
        super(piecewise_maxpool_layer, self).__init__(**kwargs)

    def build(self, input_shape):
        super(piecewise_maxpool_layer, self).build(input_shape)

    def compute_mask(self, x, mask=None):
        return None

    def max_pool_piece1(self, x):
        conv_output = x[0]
        e1 = K.cast(x[1], 'int32')
        piece = K.slice(conv_output, [0, 0], [e1, self.filter_num])
        return K.max(piece, 0)

    def max_pool_piece2(self, x):
        conv_output = x[0]
        e1 = K.cast(x[1], 'int32')
        e2 = K.cast(x[2], 'int32')
        piece = K.slice(conv_output, [e1, 0], [e2 - e1, self.filter_num])
        return K.max(piece, 0)

    def max_pool_piece3(self, x):
        conv_output = x[0]
        # ~ e2 = tf.to_int32(x[1])
        e2 = K.cast(x[1], 'int32')
        piece = K.slice(conv_output, [e2, 0], [self.seq_len - e2, self.filter_num])
        return K.max(piece, 0)

    def call(self, inputs, **kwargs):
        assert (len(inputs) == 2)
        conv_output = inputs[0]
        e1 = inputs[1][:, 1]
        e2 = inputs[1][:, 3]
        conv_piece1 = K.map_fn(self.max_pool_piece1, (conv_output, e1), dtype='float32')
        conv_piece2 = K.map_fn(self.max_pool_piece2, (conv_output, e1, e2), dtype='float32')
        conv_piece3 = K.map_fn(self.max_pool_piece3, (conv_output, e2), dtype='float32')
        return K.concatenate([conv_piece1, conv_piece2, conv_piece3])

    def compute_output_shape(self, input_shape):
        return (input_shape[0][0], input_shape[0][2] * 3)

def Squeeze(x):
    return K.squeeze(x, axis=1)

def Max(x):
    return K.max(x, axis=1)

def load_bert_tokenizer(bert_model_path, verbose=0):
    dict_path = os.path.join(bert_model_path, 'vocab.txt')
    token_dict = load_word_voc_file(dict_path, verbose)
    tokenizer = Tokenizer(token_dict)
    return token_dict, tokenizer

def load_pretrained_embedding_from_file(tcfg, word_dict):
    """
    initialize embedding matrix with pre-trained vectors from a file
    set to zeros for those not in pre-trained, is random initialization helpful?
    :param tcfg:
    :param word_dict:
    :return: word2idx, embed_matrix
    """
    # initialize word2idx
    word2idx = {word:word_dict[word] for word in word_dict if word.isupper()}
    # load pre-trained vector file
    lines = file_line2array(tcfg.word_vector_path, sepch=' ', verbose=tcfg.verbose)
    # filter the lines containing words in word_dict
    wlines = [values for values in lines if values[0] in word_dict]
    # update word dict with these words
    word2idx.update({values[0]:i+len(word2idx) for i, values in enumerate(wlines)})
    # initialize the embedding matrix
    #embed_matrix = np.zeros((len(word2idx), tcfg.embedding_dim))
    embed_matrix = np.random.normal(0, np.sqrt(3 / tcfg.embedding_dim), (len(word2idx), tcfg.embedding_dim))
    for values in wlines:
        idx = word2idx[values[0]]
        embed_matrix[idx] = np.asarray(values[1:], dtype='float32')
    #
    print('Found {:,} pretrained word vectors.'.format(len(word2idx)))
    iwords = [word for word in sorted(word_dict) if word not in word2idx]
    file_list2line(iwords, 'iwords.txt', verbose=True)
    return word2idx, embed_matrix

def load_pretrained_embedding_from_file_old(tcfg, word_dict):
    #
    lines = file_line2array(tcfg.word_vector_path, sepch=' ', verbose=tcfg.verbose)
    # filter the lines containing words in word_dict
    wlines = [values for values in lines if values[0] in word_dict]
    # build the embeddings matrix
    scope = np.sqrt(3 / tcfg.embedding_dim)
    embed_matrix = np.random.normal(0, scope, (len(word_dict), tcfg.embedding_dim))
    for values in wlines:
        idx = word_dict[values[0]]
        embed_matrix[idx] = np.asarray(values[1:], dtype='float32')
    print('Found {:,} pretrained word vectors.'.format(len(word_dict)))
    iwords = [word for word in sorted(word_dict) if word not in word_dict]
    file_list2line(iwords, 'iwords.txt', verbose=True)
    return embed_matrix

# predict on batches with different sequence length
def train_on_batch(model, doc_D, verbose=0, its=''):
    doc_gen = doc_D.__iter__()
    trange = range(doc_D.steps)
    if verbose:
        print('\nTraining {} ...'.format(its))
        trange = tqdm(trange)
    for _ in trange:
        X, Y = next(doc_gen)
        model.train_on_batch(X, Y)
    return

"""  option parameter as a list
def get_comma_separated_args(option, opt, value, parser):
    setattr(parser.values, option.dest, value.split(','))

parser = OptionParser()
parser.add_option('-f', '--foo',
                  type='string',
                  action='callback',
                  callback=get_comma_separated_args,
                  dest = foo_args_list)
"""

class OptionConfig(OptionParser):
    def __init__(self,
                 model_name='Lstm',     # model name, 'Bert', 'Lstm', 'LstmCrf'
                 epochs=3,
                 valid_ratio=0,
                 verbose=0,
                 fold_num=10,
                 fold_num_run=3,
                 max_seq_len=128,
                 batch_size=32,

                 bert_path = '',            # for BERT model
                 tokenizer=None,

                 avgmode = 'micro',
                 elabel_schema = 'SBIEO',   # set the new labeling schema
                 slabel_schema = 'IO',

                 test_ent_pred = 0,         # using predicted test entity file
                 pred_file_suff = 'ent',    # predicted file suffix

                 bld_ent_types = None,
                 bld_ent_mode = 0,
                 mark_ent_pair = 0,

                 case_sensitive = 0,    # case-sensitive for text in instances, such as SequenceLabel etc.
                 sent_simplify = 0,     # simplify sentence
                 ent_coref = 0,         # consider entity coreference
                 ent_nest = 0,          # consider nested entities
                 dbl_ent_type = 0,      # consider double typed entity
                 diff_ent_type = 0,
                 num_classes=0,

                 word_vector_path='',  # pre-trained word vector path
                 vocab_size = 0,
                 embedding_dim = 100,
                 lstm_hidden_dim = 100,
                 lstm_output_dim = 50,
                 embedding_dropout=0.2,
                 lstm_recur_dropout = 0.1,
                 lstm_output_dropout = 0.2,
                 optimizer='rmsprop'
                 ):
        super(OptionConfig,self).__init__()

        self.add_option('--model_name', dest='model_name', type='str', default=model_name, help='model name')
        self.add_option('--bertID', dest='bertID', default=('Bert' in model_name))
        self.add_option('--bert_model', dest='bert_model', default=None)
        self.add_option('--epochs', dest='epochs', default=epochs)
        self.add_option('--valid_ratio', dest='valid_ratio', default=valid_ratio)
        self.add_option('--verbose', dest='verbose', default=verbose)

        self.add_option('--fold_num', dest='fold_num', default=fold_num)
        self.add_option('--fold_num_run', dest='fold_num_run', default=fold_num_run)
        self.add_option('--max_seq_len', dest='max_seq_len', default=max_seq_len)
        self.add_option('--batch_size', dest='batch_size', default=batch_size)

        self.add_option('--bert_path', dest='bert_path', default=bert_path)
        self.add_option('--tokenizer', dest='tokenizer', default=tokenizer)
        self.add_option('--avgmode', dest='avgmode', default=avgmode)
        self.add_option('--elabel_schema', dest='elabel_schema', default=elabel_schema)
        self.add_option('--slabel_schema', dest='slabel_schema', default=slabel_schema)

        self.add_option('--test_ent_pred', dest='test_ent_pred', default=test_ent_pred)
        self.add_option('--pred_file_suff', dest='pred_file_suff', default=pred_file_suff)

        self.add_option('--bld_ent_types', dest='bld_ent_types', default=bld_ent_types) # list of entity types to be blinded
        self.add_option('--bld_ent_mode', dest='bld_ent_mode', default=bld_ent_mode)    # 0-type, 1-seq, 2-unique id
        # how to mark entity mention pair in a relation instance, 0-None, 1-#@ around, like # E1 #, @ E2 @, 2-entity type
        self.add_option('--mark_ent_pair', dest='mark_ent_pair', default=mark_ent_pair)

        self.add_option('--case_sensitive', dest='case_sensitive', default=case_sensitive)
        self.add_option('--sent_simplify', dest='sent_simplify', default=sent_simplify)
        self.add_option('--repl_spec_char', dest='repl_spec_char', default=0)
        self.add_option('--ent_coref', dest='ent_coref', default=ent_coref)
        self.add_option('--ent_nest', dest='ent_nest', default=ent_nest)
        self.add_option('--dbl_ent_type', dest='dbl_ent_type', default=dbl_ent_type)
        self.add_option('--diff_ent_type', dest='diff_ent_type', default=diff_ent_type)  # 1-relation instances with different types of entities
        self.add_option('--num_classes', dest='num_classes', default=num_classes)
        self.add_option('--test_file_id', dest='test_file_id', default=False)  # the test file in the corpus

        # non-bert network parameters
        self.add_option('--word_vector_path', dest='word_vector_path', default=word_vector_path)
        self.add_option('--vocab_size', dest='vocab_size', default=vocab_size)
        self.add_option('--embedding_dim', dest='embedding_dim', default=embedding_dim)
        self.add_option('--lstm_hidden_dim', dest='lstm_hidden_dim', default=lstm_hidden_dim)
        self.add_option('--lstm_output_dim', dest='lstm_output_dim', default=lstm_output_dim)
        self.add_option('--embedding_dropout', dest='embedding_dropout', default=embedding_dropout)
        self.add_option('--lstm_recur_dropout', dest='lstm_recur_dropout', default=lstm_recur_dropout)
        self.add_option('--lstm_output_dropout', dest='lstm_output_dropout', default=lstm_output_dropout)
        self.add_option('--optimizer', dest='optimizer', default=optimizer)

    def __str__(self):
        cfgs = ['{} = {}'.format(k, v) for k, v in sorted(self.__dict__.items(), key=lambda x:x[0])]
        return '\n'.join(cfgs)

from keras_contrib.layers.crf import CRF
from keras_contrib.losses import crf_loss
from keras_contrib.metrics import crf_viterbi_accuracy

#from sklearn_crfsuite.metrics import flat_classification_report
# https://github.com/apogiatzis/NER-BERT-conll2003/blob/master/BERT_as_Keras_Layer_Example.ipynb
def build_lstm_model(level, cfg, EMBED_MATRIX=None):
    # only x1_in is used
    x1_in = Input(shape=(None,))
    x2_in = Input(shape=(None,))
    x3_in = Input(shape=(None,))

    model = Embedding(input_dim=cfg.vocab_size, output_dim=cfg.embedding_dim, weights=[EMBED_MATRIX],
                      input_length=None, mask_zero=True)(x1_in)
    model = Dropout(cfg.embedding_dropout)(model)
    if level == 'token' or 'Att' in cfg.model_name:
        return_sequences = True
    else:
        return_sequences = False
    model = Bidirectional(LSTM(units=cfg.lstm_hidden_dim, return_sequences=return_sequences,
                               recurrent_dropout=cfg.lstm_recur_dropout))(model)  # variational biLSTM

    if 'Crf' in cfg.model_name:
        model = TimeDistributed(Dense(cfg.lstm_output_dim, activation="relu"))(model)  # a dense layer as suggested by neuralNer
        output = CRF(cfg.num_classes[0])(model)  # CRF layer
    elif level == 'token':
        output = TimeDistributed(Dense(cfg.num_classes[0], activation="softmax"))(model)  # a dense layer as suggested by neuralNer
    else:
        if 'Att' in cfg.model_name:
            model = SelfAttention()(model)
        output = Dense(cfg.num_classes[0], activation="softmax")(model)

    # model = Dropout(cfg.lstm_output_dropout)(model)
    model = Model([x1_in, x2_in, x3_in], [output])
    if 'Crf' in cfg.model_name:
        model.compile(optimizer=cfg.optimizer, loss=crf_loss, metrics=[crf_viterbi_accuracy])
    else:
        model.compile(optimizer=cfg.optimizer, loss='categorical_crossentropy', metrics=['categorical_accuracy'])
    return model

def build_cnn_model(level, cfg, EMBED_MATRIX=None):
    x1_in = Input(shape=(None,))
    x2_in = Input(shape=(None,))
    x3_in = Input(shape=(None,))

    word_embed = Embedding(input_dim=cfg.vocab_size, output_dim=cfg.embedding_dim, weights=[EMBED_MATRIX],
                           input_length=None, mask_zero=False)(x1_in)

    model = Conv1D(filters=cfg.embedding_dim, kernel_size=3)(word_embed)

    if 'Piece' in cfg.model_name:
        model = piecewise_maxpool_layer(cfg.embedding_dim, cfg.max_seq_len-2)([model, x3_in])
    else:
        #model = MaxPooling1D(pool_size=cfg.max_seq_len-2)(model)
        #model = Lambda(Squeeze)(model)
        model = Lambda(Max)(model)

    output = Dense(cfg.num_classes[0], activation='softmax')(model)

    model = Model([x1_in, x2_in, x3_in], [output])
    model.compile(optimizer=cfg.optimizer, loss='categorical_crossentropy', metrics=['categorical_accuracy'])
    return model

def build_simple_lstm_crf_model(level, model_name, num_classes, vocab_size, EMBED_MATRIX=None, optimizer='rmsprop'):
    """
    used out of the project
    :param level: 'sent'-sentence, 'token'-token
    :param model_name:
    :param num_classes:
    :param vocab_size:
    :param EMBED_MATRIX:
    :param optimizer:
    :return:
    """
    # only x1_in is used
    input = Input(shape=(None,))

    model = Embedding(input_dim=vocab_size, output_dim=100, weights=[EMBED_MATRIX],
                      input_length=None, mask_zero=False)(input)
    model = Dropout(0.2)(model)
    if level == 'sent' and 'Att' not in model_name:
        return_sequences = False
    else:
        return_sequences = False
    model = Bidirectional(LSTM(units=100, return_sequences=return_sequences, recurrent_dropout=0.1))(model)  # variational biLSTM

    if 'Crf' in model_name:
        model = TimeDistributed(Dense(50, activation="relu"))(model)  # a dense layer as suggested by neuralNer
        output = CRF(num_classes)(model)  # CRF layer
    else:
        output = TimeDistributed(Dense(num_classes, activation="softmax"))(model)  # a dense layer as suggested by neuralNer

    model = Model(input, output)
    if 'Crf' in model_name:
        model.compile(optimizer=optimizer, loss=crf_loss, metrics=[crf_viterbi_accuracy])
    else:
        model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['categorical_accuracy'])
    return model

def create_bert_sequence_labeling_model(bert_model_path, cfg, verbose=0):
    if verbose:  print('\nCreating {} model ...'.format(cfg.model_name))
    config_path = os.path.join(bert_model_path, 'bert_config.json')
    checkpoint_path = os.path.join(bert_model_path, 'bert_model.ckpt')

    bert_model = load_trained_model_from_checkpoint(config_path, checkpoint_path, seq_len=None)

    for l in bert_model.layers: l.trainable = True

    x1_in = Input(shape=(None,))
    x2_in = Input(shape=(None,))

    model = bert_model([x1_in, x2_in])
    if 'Lstm' in cfg.model_name:
        model = Bidirectional(LSTM(units=cfg.lstm_hidden_dim, return_sequences=True,
                                   recurrent_dropout=cfg.recurrent_dropout))(model)  # variational biLSTM

    if 'Crf' in cfg.model_name:
        model = TimeDistributed(Dense(cfg.lstm_output_dim, activation="relu"))(model)  # a dense layer as suggested by neuralNer
        output = CRF(cfg.tag_size)(model)  # CRF layer
    else:   # Bert and BertLstm
        output = TimeDistributed(Dense(cfg.tag_size, activation="softmax"))(model)

    model = Model([x1_in, x2_in], [output])
    if 'Crf' in cfg.model_name:
        model.compile(loss=crf_loss, optimizer=Adam(1e-5), metrics=[crf_viterbi_accuracy])
    else:
        model.compile(loss='categorical_crossentropy', optimizer=Adam(1e-5), metrics=['categorical_accuracy'])
    model.summary()
    return model

def load_bert_model(bert_path, verbose=0):
    if verbose:
        print('\nLoading the BERT model from {} ...'.format(bert_path))
    config_path = os.path.join(bert_path, 'bert_config.json')
    checkpoint_path = os.path.join(bert_path, 'bert_model.ckpt')
    bert_model = load_trained_model_from_checkpoint(config_path, checkpoint_path, seq_len=None)
    return bert_model

# level: sent, token
def create_bert_classification_model(level='sent', tcfg=None):

    for l in tcfg.bert_model.layers: l.trainable = True

    x1_in = Input(shape=(None,))
    x2_in = Input(shape=(None,))
    x3_in = Input(shape=(None,))

    x = tcfg.bert_model([x1_in, x2_in])
    if level == 'token':
        if 'Crf' in tcfg.model_name:
            x = TimeDistributed(Dense(tcfg.lstm_output_dim, activation="relu"))(x)  # a dense layer as suggested by neuralNer
            output = CRF(tcfg.num_classes[0])(x)  # CRF layer
        else:
            output = TimeDistributed(Dense(tcfg.num_classes[0], activation="softmax"))(x)
    else:
        x = Lambda(lambda x: x[:, 0])(x)
        output = Dense(tcfg.num_classes[0], activation='softmax')(x)

    model = Model([x1_in, x2_in, x3_in], [output])
    if 'Crf' in tcfg.model_name:
        model.compile(loss=crf_loss, optimizer=Adam(1e-5), metrics=[crf_viterbi_accuracy])
    else:
        model.compile(loss='categorical_crossentropy', optimizer=Adam(1e-5), metrics=['categorical_accuracy'])
    return model

# level: sent, token
def create_bert_ve_arg_classification_model(bert_model=None, num_classes=0):

    for l in bert_model.layers: l.trainable = True

    x1_in = Input(shape=(None,))
    x2_in = Input(shape=(None,))
    x3_in = Input(shape=(None,))

    # 1-Bind, 2-GeEx, 3-Loca, 4-NeRe, 5-Phos, 6-PoRe, 7-PrCa, 8-Regu, 9-Tran
    x = bert_model([x1_in, x2_in])
    output = TimeDistributed(Dense(num_classes, activation="softmax"))(x)

    model = Model([x1_in, x2_in, x3_in], [output])
    model.compile(loss='categorical_crossentropy', optimizer=Adam(1e-5), metrics=['categorical_accuracy'])
    return model


"""
2021 Day 15
https://adventofcode.com/2021/day/15
"""

from heapq import heappop, heappush
from typing import Tuple
import aocd  # type: ignore
import numpy as np

def read_map(text: str) -> np.ndarray:
    return np.array([list(line) for line in text.split("\n")], int)

MAX_INT = np.iinfo(np.int64).max

def lowest_risk_path(risk_map: np.ndarray, initial: Tuple[int, int] = (0, 0)) -> float:
    max_y = risk_map.shape[0] - 1
    max_x = risk_map.shape[1] - 1
    risks = np.ones_like(risk_map) * MAX_INT
    risks[initial] = 0
    queue = [(0, initial)]

    while queue:
        current_risk, current_pos = heappop(queue)

        if current_pos == (max_y, max_x):
            return current_risk

        if current_risk > risks[current_pos]:
            continue

        y, x = current_pos
        for dy, dx in ((0, 1), (0, -1), (1, 0), (-1, 0)):
            ny, nx = y + dy, x + dx
            if 0 <= ny <= max_y and 0 <= nx <= max_x:
                risk = current_risk + risk_map[ny, nx]
                if risk < risks[ny, nx]:
                    risks[ny, nx] = risk
                    heappush(queue, (risk, (ny, nx)))

    return -1

def increment(risk_map: np.ndarray, increment_by: int) -> np.ndarray:
    incremented = risk_map + increment_by
    while (incremented > 9).any():
        incremented[incremented > 9] -= 9
    return incremented

def extend_risk_map(risk_map: np.ndarray) -> np.ndarray:
    increments = [
        [0, 1, 2, 3, 4],
        [1, 2, 3, 4, 5],
        [2, 3, 4, 5, 6],
        [3, 4, 5, 6, 7],
        [4, 5, 6, 7, 8],
    ]
    return np.block(
        [[increment(risk_map, value) for value in row] for row in increments]
    )

def test_part1() -> None:
    """
    Examples for Part 1.
    """
    risk_map = np.array(
        [
            [1, 1, 6, 3, 7, 5, 1, 7, 4, 2],
            [1, 3, 8, 1, 3, 7, 3, 6, 7, 2],
            [2, 1, 3, 6, 5, 1, 1, 3, 2, 8],
            [3, 6, 9, 4, 9, 3, 1, 5, 6, 9],
            [7, 4, 6, 3, 4, 1, 7, 1, 1, 1],
            [1, 3, 1, 9, 1, 2, 8, 1, 3, 7],
            [1, 3, 5, 9, 9, 1, 2, 4, 2, 1],
            [3, 1, 2, 5, 4, 2, 1, 6, 3, 9],
            [1, 2, 9, 3, 1, 3, 8, 5, 2, 1],
            [2, 3, 1, 1, 9, 4, 4, 5, 8, 1],
        ]
    )
    assert (
        read_map(
            "\n".join(
                (
                    "1163751742",
                    "1381373672",
                    "2136511328",
                    "3694931569",
                    "7463417111",
                    "1319128137",
                    "1359912421",
                    "3125421639",
                    "1293138521",
                    "2311944581",
                )
            )
        )
        == risk_map
    ).all()
    assert lowest_risk_path(risk_map) == 40

def test_part2() -> None:
    """
    Examples for Part 2.
    """
    risk_map = np.array(
        [
            [1, 1, 6, 3, 7, 5, 1, 7, 4, 2],
            [1, 3, 8, 1, 3, 7, 3, 6, 7, 2],
            [2, 1, 3, 6, 5, 1, 1, 3, 2, 8],
            [3, 6, 9, 4, 9, 3, 1, 5, 6, 9],
            [7, 4, 6, 3, 4, 1, 7, 1, 1, 1],
            [1, 3, 1, 9, 1, 2, 8, 1, 3, 7],
            [1, 3, 5, 9, 9, 1, 2, 4, 2, 1],
            [3, 1, 2, 5, 4, 2, 1, 6, 3, 9],
            [1, 2, 9, 3, 1, 3, 8, 5, 2, 1],
            [2, 3, 1, 1, 9, 4, 4, 5, 8, 1],
        ]
    )
    extended = extend_risk_map(risk_map)
    assert extended.shape == (50, 50)
    assert extended[1, 2] == 8
    assert extended[1, 12] == 9
    assert extended[1, 22] == 1
    assert extended[1, 32] == 2
    assert extended[1, 42] == 3
    assert extended[11, 2] == 9
    assert extended[11, 12] == 1
    assert extended[11, 22] == 2
    assert extended[11, 32] == 3
    assert extended[11, 42] == 4
    assert extended[41, 42] == 7
    assert lowest_risk_path(extended) == 315

def main() -> None:
    """
    Calculate and output the solutions based on the real puzzle input.
    """
    data = aocd.get_data(year=2021, day=15)

    risk_map = read_map(data)
    print(f"Part 1: {lowest_risk_path(risk_map)}")

    extended = extend_risk_map(risk_map)
    print(f"Part 2: {lowest_risk_path(extended)}")

if __name__ == "__main__":
    main()


# Run with: python simple_server.py
import ujson

from tornado import ioloop, web

class MainHandler(web.RequestHandler):
    def get(self):
        self.write(ujson.dumps({"test": True}))

app = web.Application(
    [(r"/", MainHandler)],
    debug=False,
    compress_response=False,
    static_hash_cache=True,
)

app.listen(8000)
ioloop.IOLoop.current().start()


import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init

def kaiming_init(m):
    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.Linear)):
        init.kaiming_normal_(m.weight)
        if m.bias is not None:
            m.bias.data.fill_(0)

def normal_init(m):
    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.Linear)):
        m.weight.data.normal_(mean=0, std=0.02)
        if m.bias.data is not None:
            m.bias.data.zero_()

def reparametrize(mu, logvar):
    std = logvar.div(2).exp()
    eps = std.data.new(std.size()).normal_()
    return mu + std*eps

class View(nn.Module):
    def __init__(self, size):
        super(View, self).__init__()
        self.size = size
    def forward(self, tensor):
        return tensor.view(self.size)

class Encoder(nn.Module):
    def __init__(self, c_dim=10, nc=3, infodistil_mode=False):
        super(Encoder, self).__init__()
        self.c_dim = c_dim
        self.nc = nc
        self.infodistil_mode = infodistil_mode
        self.layer = nn.Sequential(
            nn.Conv2d(nc, 32, 4, 2, 1),          # B,  32, 32, 32
            nn.ReLU(True),
            nn.Conv2d(32, 32, 4, 2, 1),          # B,  32, 16, 16
            nn.ReLU(True),
            nn.Conv2d(32, 64, 4, 2, 1),          # B,  64,  8,  8
            nn.ReLU(True),
            nn.Conv2d(64, 64, 4, 2, 1),          # B,  64,  4,  4
            nn.ReLU(True),
            nn.Conv2d(64, 256, 4, 1),            # B, 256,  1,  1
            nn.ReLU(True),
            View((-1, 256*1*1)),                 # B, 256
            nn.Linear(256, c_dim*2),             # B, c_dim*2
        )

    def forward(self, x):
        if self.infodistil_mode:
            x = x.add(1).div(2)
            if (x.size(2) > 64) or (x.size(3) > 64):
                x = F.adaptive_avg_pool2d(x, (64, 64))

        h = self.layer(x)
        return h

class Decoder(nn.Module):
    def __init__(self, c_dim=10, nc=3):
        super(Decoder, self).__init__()
        self.c_dim = c_dim
        self.nc = nc
        self.layer = nn.Sequential(
            nn.Linear(c_dim, 256),               # B, 256
            View((-1, 256, 1, 1)),               # B, 256,  1,  1
            nn.ReLU(True),
            nn.ConvTranspose2d(256, 64, 4),      # B,  64,  4,  4
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 64, 4, 2, 1), # B,  64,  8,  8
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 32, 4, 2, 1), # B,  32, 16, 16
            nn.ReLU(True),
            nn.ConvTranspose2d(32, 32, 4, 2, 1), # B,  32, 32, 32
            nn.ReLU(True),
            nn.ConvTranspose2d(32, nc, 4, 2, 1),  # B, nc, 64, 64
        )

    def forward(self, c):
        x = self.layer(c)
        return x

class Discriminator(nn.Module):
    def __init__(self, z_dim, size=64, nc=3):
        super(Discriminator, self).__init__()
        self.nc = nc
        self.layer = nn.Sequential(
            nn.Conv2d(nc, 32, 4, 2, 1),          # B,  32, 32, 32
            nn.ReLU(True),
            nn.Conv2d(32, 32, 4, 2, 1),          # B,  32, 16, 16
            nn.ReLU(True),
            nn.Conv2d(32, 64, 4, 2, 1),          # B,  64,  8,  8
            nn.ReLU(True),
            nn.Conv2d(64, 64, 4, 2, 1),          # B,  64,  4,  4
            nn.ReLU(True),
            nn.Conv2d(64, 256, 4, 1),            # B, 256,  1,  1
            nn.ReLU(True),
            View((-1, 256*1*1)),                 # B, 256
            nn.Linear(256, 1),             # B, 1
        )

    def forward(self, x):
        out = self.layer(x)
        return out

class Generator(nn.Module):
    def __init__(self, z_dim=10, size=64, nc=3, **kwarngs):
        super(Generator, self).__init__()
        self.z_dim = z_dim
        self.nc = nc
        self.layer = nn.Sequential(
            nn.Linear(z_dim, 256),               # B, 256
            View((-1, 256, 1, 1)),               # B, 256,  1,  1
            nn.ReLU(True),
            nn.ConvTranspose2d(256, 64, 4),      # B,  64,  4,  4
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 64, 4, 2, 1), # B,  64,  8,  8
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 32, 4, 2, 1), # B,  32, 16, 16
            nn.ReLU(True),
            nn.ConvTranspose2d(32, 32, 4, 2, 1), # B,  32, 32, 32
            nn.ReLU(True),
            nn.ConvTranspose2d(32, nc, 4, 2, 1),  # B, nc, 64, 64
            nn.Tanh()
        )

    def forward(self, c):
        x = self.layer(c)
        return x

class BetaVAE_H(nn.Module):
    """Model proposed in original beta-VAE paper(Higgins et al, ICLR, 2017)."""

    def __init__(self, c_dim=10, nc=3, infodistil_mode=False):
        super(BetaVAE_H, self).__init__()
        self.c_dim = c_dim
        self.nc = nc
        self.encoder = Encoder(c_dim, nc, infodistil_mode)
        self.decoder = Decoder(c_dim, nc)
        self.apply(normal_init)

    def forward(self, x, c, encode_only, decode_only):
        if encode_only:
            c, mu, logvar = self._encode(x)
            return c, mu, logvar
        elif decode_only:
            x_recon = self._decode(c)
            return x_recon
        else:
            c, mu, logvar = self._encode(x)
            x_recon = self._decode(c)
            return x_recon, c, mu, logvar

    def __call__(self, x=None, c=None, encode_only=False, decode_only=False):
        return self.forward(x, c, encode_only, decode_only)

    def _encode(self, x):
        distributions = self.encoder(x)
        mu = distributions[:, :self.c_dim]
        logvar = distributions[:, self.c_dim:]
        c = reparametrize(mu, logvar)
        return c, mu, logvar

    def _decode(self, c):
        return self.decoder(c)


from datetime import date, timedelta

def date_convert(value):
    value_date = map(int, value.split('-'))
    value_date = date(*value_date)
    return value_date

def time_convert(value):
    minutes = value % 60
    hours = value // 60
    return '{}:{:0>2}'.format(hours, minutes)

if __name__ == '__main__':
    print(time_convert(130))


import sys
from pathlib import Path
from parse import Parse
from codewriter import CodeWriter
from cmdtable import CmdTable

def write_vm1(inputfile, in_filename, ctabel, cw):
    par = Parse(inputfile, ctabel)
    cw.setInputname(in_filename)

    while par.hasMoreCommands():
        par.advance()
        ctype = par.cmdType()
        if ctype == "C_ARITHMETIC":
            cw.writeArithmetic(par.arg1())
        elif ctype == "C_PUSH" or ctype == "C_POP":
            cw.writePushPop(ctype, par.arg1(), int(par.arg2()))
        elif ctype == "C_LABEL":
            cw.writeLabel(par.arg1())
        elif ctype == "C_IF":
            cw.writeIf(par.arg1())
        elif ctype == "C_GOTO":
            cw.writeGoto(par.arg1())
        elif ctype == "C_FUNCTION":
            cw.writeFunction(par.arg1(), int(par.arg2()))
        elif ctype == "C_RETURN":
            cw.writeReturn()
        elif ctype == "C_CALL":
            cw.writeCall(par.arg1(), int(par.arg2()))

    #cw.getarith_jump()

def main():
    p=Path(sys.argv[1])
    ctabel = CmdTable()
    out_filename= p.stem
    inputfile_list = {}
    if p.is_file():
        outputfile = str(p.parent / (out_filename + ".asm"))
        inputfile_list[str(p)] = p.stem
    elif p.is_dir():
        outputfile = p / (out_filename + ".asm")
        for pfile in p.glob('*.vm'):
            inputfile_list[str(pfile)] = pfile.stem

    fout = open(outputfile, mode='w')
    cw = CodeWriter(fout, ctabel)
    if p.is_dir():
        cw.writeBoot()

    for inf, infname in inputfile_list.items():
        write_vm1(inf, infname, ctabel, cw)
    fout.close()

    # debug: write line number
    #newfile = p / (out_filename + ".txt")
    #print(newfile)
    #with open(newfile, 'w') as file_linenum:
    #    line_count = 0
    #    with open(outputfile) as file_asm:
    #        for line in file_asm:
    #            if ("//" not in line) and ("(" not in line):
    #                line_count += 1
    #                line = str(line_count) + " " + line
    #            file_linenum.write(line)

if __name__ == "__main__":
    main()


import itertools

import numpy as np
import plotly.graph_objects as go
import pandas as pd
import pathlib
import plotly.express as px
from plotly.subplots import make_subplots

PATH = pathlib.Path(__file__).parent
DATA_PATH = (PATH.joinpath("data")).resolve()
encoding_names = {'initial_configuration': "O'SFS + CL
(Initial encoding)",
                  "no_output_before_pop":  "O'SFS + CL + CR",
                  'at_most':  "O'SFS + CL + CU",
                  'pushed_once':  "O'SFS + CL + CN",
                  "alternative_gas_model": "OSFS + CL",
                  "final_encoding": "Final encoding",
                  "60s": "60s", "30s": "30s", "15s": "15s", "10s": "10s", "1s": "1s",
                  "no_output_before_pop_at_most": "O'SFS + CL + CR + CU",
                  "no_output_before_pop_pushed_once": "O'SFS + CL + CR + CN",
                  "no_output_before_pop_at_most_pushed_once": "O'SFS + CL "
                                                              "+ CR + CU + CN",
                  "final_setup": "Final setup"}

optimality_names = {'already_optimal': "Already optimal", 'discovered_optimal': "Discovered optimal",
                    'non_optimal_with_less_gas': "Non optimal
 with less gas",
                    'non_optimal_with_same_gas': "Non optimal
 with same gas",
                    'no_solution_found': "No solution found"}

encoding_names_abbreviated = {'initial_configuration': "initial", "no_output_before_pop": "restrict POP",
                              'at_most': "at most", 'pushed_once': "values pushed",
                              "alternative_gas_model": "alt. gas model", "60s": "60s", "30s": "30s", "15s": "15s",
                              "10s": "10s", "1s": "1s", "final_encoding": "final"}

analyzed_parameters_names = {'saved_gas': "Gas saved", 'time': "Time in min"}

solver_name_abbreviated = {"combined": "port", "z3": "z3", "oms": "oms", "barcelogic": "bclg"}

solver_name = {"combined": "Portfolio", "z3": "Z3", "oms": "OMS", "barcelogic": "Barcelogic"}

def plot_time(folder_name, encodings):
    fig = go.Figure()
    for encoding in encodings:
        times = np.empty(0, dtype=float)
        labels = []
        for name in folder_name:
            csv_name = str(DATA_PATH) + "/" + encoding + "_" + name + ".csv"
            df = pd.read_csv(csv_name)
            arr = df['time'].to_numpy() / 60
            times = np.append(times, arr)
            labels.extend([solver_name[name]] * len(arr))
        fig.add_trace(go.Box(y=times, x=labels, name=encoding_names[encoding]))
    fig.update_layout(
        yaxis_title='Times per contract (minutes)',
        boxmode='group'  # group together boxes of the different traces for each value of x
    )
    return fig

def plot_gas(folder_name, encodings):
    fig = go.Figure()
    for encoding in encodings:
        times = np.empty(0, dtype=float)
        labels = []
        for name in folder_name:
            csv_name = str(DATA_PATH) + "/" + encoding + "_" + name + ".csv"
            df = pd.read_csv(csv_name)
            arr = df['saved_gas'].to_numpy()
            times = np.append(times, arr)
            labels.extend([solver_name[name]] * len(arr))
        fig.add_trace(go.Box(y=times, x=labels, name=encoding_names[encoding]))
    fig.update_layout(
        yaxis_title='Saved gas per contract',
        boxmode='group'  # group together boxes of the different traces for each value of x
    )
    return fig

def plot_statistics(folder_name, encodings):
    fig = go.Figure()
    statistics = ['already_optimal', 'discovered_optimal', 'non_optimal_with_less_gas', 'non_optimal_with_same_gas',
                  'no_solution_found']
    for statistic in statistics:
        labels_x = []
        labels_y = []
        results = []
        for name in folder_name:
            for encoding in encodings:
                csv_name = str(DATA_PATH) + "/" + encoding + "_" + name + ".csv"
                df = pd.read_csv(csv_name)
                total_sum = 0
                for other_statistics in statistics:
                    total_sum += df[other_statistics].sum()
                labels_x.append(solver_name_abbreviated[name])
                labels_y.append(encoding_names_abbreviated[encoding])
                results.append((df[statistic].sum() * 100) / total_sum)
        fig.add_trace(go.Bar(y=results, x=[labels_x, labels_y], name=optimality_names[statistic]))
    fig.update_layout(
        yaxis_title='Comparison in outputs',
        barmode='stack',
    )
    return fig

def select_comparison(df, comparison_category):
    if comparison_category == "init_progr_len":
        return df[comparison_category].to_numpy()
    elif comparison_category == "initial_size_relation":
        return df[comparison_category].to_numpy()
    elif comparison_category == "number_of_necessary_push":
        return df[comparison_category].to_numpy()
    elif comparison_category == "number_of_necessary_uninterpreted_instructions":
        return df[comparison_category].to_numpy()
    elif comparison_category == "push_per_initial":
        return df["number_of_necessary_push"].to_numpy() / df["init_progr_len"].to_numpy()
    elif comparison_category == "uninterpreted_per_initial":
        return df["number_of_necessary_uninterpreted_instructions"].to_numpy() / df["init_progr_len"].to_numpy()
    elif comparison_category == "push_per_expected":
        expected_size = np.array(list(map(lambda x: 1 if x == 0 else x, df['inferred_size'])))
        return df["number_of_necessary_push"].to_numpy() / expected_size
    elif comparison_category == "uninterpreted_per_expected":
        expected_size = np.array(list(map(lambda x: 1 if x == 0 else x, df['inferred_size'])))
        return df["number_of_necessary_uninterpreted_instructions"].to_numpy() / expected_size

def plot_comparison(cat1, cat2, relation):
    csv_name1 = str(DATA_PATH) + "/" + "comparison_" + cat1 + "_" + cat2 + ".csv"
    csv_name2 = str(DATA_PATH) + "/" + "comparison_" + cat2 + "_" + cat1 + ".csv"
    y1 = select_comparison(pd.read_csv(csv_name1), relation)
    y2 = select_comparison(pd.read_csv(csv_name2), relation)
    fig = go.Figure()
    fig.add_trace(go.Box(y=y1, name="Default encoding
works better", boxpoints='all', marker_size=3))
    fig.add_trace(go.Box(y=y2, name="Selected encoding
works better", boxpoints='all', marker_size=3))
    fig.update_layout(yaxis_title="Comparison between encodings")
    return fig

def plot_configuration_comparison(category_comparison):
    csv_name = str(DATA_PATH) + "/" + category_comparison + "_parameter_comparison.csv"
    df = pd.read_csv(csv_name)
    x = [encoding_names[encoding] for encoding in df['name'].to_list()]
    y = df['time'].to_list()
    fig = go.Figure(data=[go.Bar(x=x, y=y)])
    fig.update_layout(yaxis_title="Total time in minutes")
    return fig

def plot_statistics_pie_chart(solver):
    syrup_csv_name = str(DATA_PATH) + "/final_setup_" + solver + ".csv"
    cav_csv_name = str(DATA_PATH) + "/CAV_" + solver + ".csv"
    labels = ['already_optimal', 'discovered_optimal', 'non_optimal_with_less_gas',
              'non_optimal_with_same_gas', 'no_solution_found']
    labels_to_desplay = [optimality_names[name] for name in labels]
    cav_df = pd.read_csv(cav_csv_name).sum()
    syrup_df = pd.read_csv(syrup_csv_name).sum()

    cav_values = [cav_df[label] for label in labels]
    syrup_values = [syrup_df[label] for label in labels]

    fig = make_subplots(rows=1, cols=2, specs=[[{'type': 'domain'}, {'type': 'domain'}]])
    fig.add_trace(go.Pie(labels=labels_to_desplay, values=cav_values, name="CAV'20 Setup"),
                  1, 1)
    fig.add_trace(go.Pie(labels=labels_to_desplay, values=syrup_values, name="New setup"),
                  1, 2)

    # Use `hole` to create a donut-like pie chart
    fig.update_traces(hole=.4, hoverinfo="label+percent+name")

    fig.update_layout(
        title_text="Comparison between Syrup 1.0 setup (15 min) vs Syrup 2.0 setup (10 s)",
        # Add annotations in the center of the donut pies.
        annotations=[dict(text='syrup 1.0', x=0.14, y=0.5, font_size=15, showarrow=False),
                     dict(text='syrup 2.0', x=0.85, y=0.5, font_size=15, showarrow=False)])
    return fig

def plot_bar_comparison(solver, category_name):
    syrup_csv_name = str(DATA_PATH) + "/final_setup_" + solver + ".csv"
    cav_csv_name = str(DATA_PATH) + "/CAV_" + solver + ".csv"
    cav_rows = pd.read_csv(cav_csv_name).to_dict('records')
    syrup_rows = pd.read_csv(syrup_csv_name).to_dict('records')
    labels = list(range(len(syrup_rows)))
    cav_values = []
    syrup_values = []
    for i, cav_row in enumerate(cav_rows):
        block_name = cav_row['name']
        syrup_row = list(filter(lambda row: row['name'] == block_name, syrup_rows))[0]
        if category_name == "time":
            cav_values.append(cav_row[category_name] / 60)
            syrup_values.append(syrup_row[category_name] / 60)
        else:
            cav_values.append(cav_row[category_name])
            syrup_values.append(syrup_row[category_name])
    fig = go.Figure(data=[go.Bar(name="syrup 1.0", x=labels, y=cav_values),
                          go.Bar(name='syrup 2.0', x=labels, y=syrup_values)])
    fig.update_layout(barmode='group', font_size=16, yaxis_title=analyzed_parameters_names[category_name] + ' per contract',)
    if category_name == "time":
        fig.update_layout(yaxis_type="log")
    return fig


"""Parameter selector variables"""

import os
import math

import pyomo.environ as pyo

from utils.data import get_case_data

def define_sets(m, data):
    """Define sets"""

    # Nodes
    m.OMEGA_N = pyo.Set(initialize=data['S_NODES'])

    # Generators
    m.OMEGA_G = pyo.Set(initialize=data['S_GENERATORS'])

    # AC edges
    m.OMEGA_NM = pyo.Set(initialize=data['S_AC_EDGES'])

    # Sets of branches for which aggregate AC interconnector limits are defined
    m.OMEGA_J = pyo.Set(initialize=data['S_AC_AGGREGATED_EDGES'])

    # HVDC links
    m.OMEGA_H = pyo.Set(initialize=data['S_HVDC_EDGES'])

    # Operating scenarios
    m.OMEGA_S = pyo.Set(initialize=data['S_SCENARIOS'])

    # NEM regions
    m.OMEGA_R = pyo.Set(initialize=data['S_REGIONS'])

    # Branches which constitute AC interconnectors in the network
    m.OMEGA_L = pyo.Set(initialize=data['S_INTERCONNECTORS'])

    # Set of integers used to select discretised baseline
    m.OMEGA_U = pyo.Set(initialize=data['S_BINARY_EXPANSION_INTEGERS'])

    return m

def define_network_parameters(m, data):
    """Define network parameters"""

    # Max voltage angle difference between connected nodes
    m.P_NETWORK_VOLTAGE_ANGLE_MAX_DIFFERENCE = pyo.Param(initialize=float(math.pi / 2), mutable=True)

    # Branch susceptance matrix elements
    m.P_NETWORK_BRANCH_SUSCEPTANCE = pyo.Param(m.OMEGA_NM, initialize=data['P_NETWORK_BRANCH_SUSCEPTANCE'],
                                               mutable=True)

    # HVDC power flow limits
    m.P_NETWORK_HVDC_FORWARD_LIMIT = pyo.Param(m.OMEGA_H, initialize=data['P_NETWORK_HVDC_FORWARD_LIMIT'], mutable=True)
    m.P_NETWORK_HVDC_REVERSE_LIMIT = pyo.Param(m.OMEGA_H, initialize=data['P_NETWORK_HVDC_REVERSE_LIMIT'], mutable=True)

    # Reference node indicator
    m.P_NETWORK_REFERENCE_NODE_INDICATOR = pyo.Param(m.OMEGA_N, initialize=data['P_NETWORK_REFERENCE_NODE_INDICATOR'])

    # HVDC incidence matrix
    m.P_NETWORK_HVDC_INCIDENCE_MAT = pyo.Param(m.OMEGA_N, m.OMEGA_H, initialize=data['P_NETWORK_HVDC_INCIDENCE_MAT'])

    # AC interconnector incidence matrix
    m.P_NETWORK_INTERCONNECTOR_INCIDENCE_MAT = pyo.Param(m.OMEGA_N, m.OMEGA_L,
                                                         initialize=data['P_NETWORK_INTERCONNECTOR_INCIDENCE_MAT'])

    # Aggregate AC interconnector flow limits
    m.P_NETWORK_INTERCONNECTOR_LIMIT = pyo.Param(m.OMEGA_J, initialize=data['P_NETWORK_INTERCONNECTOR_LIMIT'],
                                                 mutable=True)

    # Fixed nodal power injections
    m.P_NETWORK_FIXED_INJECTION = pyo.Param(m.OMEGA_S, m.OMEGA_N, initialize=data['P_NETWORK_FIXED_INJECTION'],
                                            mutable=True)

    # Network demand
    m.P_NETWORK_DEMAND = pyo.Param(m.OMEGA_S, m.OMEGA_N, initialize=data['P_NETWORK_DEMAND'], mutable=True)

    # Proportion of total demand consumed in each region
    m.P_NETWORK_REGION_DEMAND_PROPORTION = pyo.Param(m.OMEGA_S, m.OMEGA_R,
                                                     initialize=data['P_NETWORK_REGION_DEMAND_PROPORTION'])

    # Network graph describing connections between nodes
    m.P_NETWORK_GRAPH = pyo.Param(m.OMEGA_N, initialize=data['P_NETWORK_GRAPH'], within=pyo.Any)

    # Branches associated with each interconnector
    m.P_NETWORK_INTERCONNECTOR_BRANCH_ID_MAP = pyo.Param(m.OMEGA_J,
                                                         initialize=data['P_NETWORK_INTERCONNECTOR_BRANCH_ID_MAP'],
                                                         within=pyo.Any)

    # Mapping between branch IDs and nodes
    m.P_NETWORK_INTERCONNECTOR_BRANCH_NODES_MAP = pyo.Param(
        m.OMEGA_L, initialize=data['P_NETWORK_INTERCONNECTOR_BRANCH_NODES_MAP'], within=pyo.Any)

    # Mapping between HVDC link ID and its respective 'from' and 'to' nodes
    m.P_NETWORK_HVDC_FROM_NODE = pyo.Param(m.OMEGA_H, initialize=data['P_NETWORK_HVDC_FROM_NODE'])
    m.P_NETWORK_HVDC_TO_NODE = pyo.Param(m.OMEGA_H, initialize=data['P_NETWORK_HVDC_TO_NODE'])

    # List of generators assigned to each node
    m.P_NETWORK_NODE_GENERATOR_MAP = pyo.Param(m.OMEGA_N, initialize=data['P_NETWORK_NODE_GENERATOR_MAP'],
                                               within=pyo.Any)

    # Base power
    m.P_NETWORK_BASE_POWER = pyo.Param(initialize=100)

    # Regional reference node IDs for each region
    m.P_REGION_RRN_NODE_MAP = pyo.Param(m.OMEGA_R, initialize=data['P_REGION_RRN_NODE_MAP'])

    return m

def define_scenario_parameters(m, data):
    """Define scenario parameters"""

    # Relative scenario duration
    m.P_SCENARIO_DURATION = pyo.Param(m.OMEGA_S, initialize=data['P_SCENARIO_DURATION'])

    return m

def define_generator_parameters(m, data):
    """Define generator parameters"""

    # Maximum generator output
    m.P_GENERATOR_MAX_OUTPUT = pyo.Param(m.OMEGA_G, initialize=data['P_GENERATOR_MAX_OUTPUT'], mutable=True)

    # Minimum generator output (set to 0)
    m.P_GENERATOR_MIN_OUTPUT = pyo.Param(m.OMEGA_G, initialize=data['P_GENERATOR_MIN_OUTPUT'])

    # Generator short-run marginal costs
    m.P_GENERATOR_SRMC = pyo.Param(m.OMEGA_G, initialize=data['P_GENERATOR_SRMC'], mutable=True)

    # Generator emissions intensities
    m.P_GENERATOR_EMISSIONS_INTENSITY = pyo.Param(m.OMEGA_G, initialize=data['P_GENERATOR_EMISSIONS_INTENSITY'])

    # Node to which generator is assigned
    m.P_GENERATOR_NODE = pyo.Param(m.OMEGA_G, initialize=data['P_GENERATOR_NODE'])

    return m

def define_policy_parameters(m, data):
    """Define policy parameters"""

    # Emissions intensity baseline (fixed)
    m.P_POLICY_FIXED_BASELINE = pyo.Param(initialize=data['P_POLICY_FIXED_BASELINE'], mutable=True)

    # Permit price target
    m.P_POLICY_PERMIT_PRICE_TARGET = pyo.Param(initialize=data['P_POLICY_PERMIT_PRICE_TARGET'], mutable=True)

    # Weighted RRN price target
    m.P_POLICY_WEIGHTED_RRN_PRICE_TARGET = pyo.Param(initialize=data['P_POLICY_WEIGHTED_RRN_PRICE_TARGET'],
                                                     mutable=True)

    return m

def define_big_m_parameters(m):
    """Define Big-M parameters"""

    def m_11_rule(m, g):
        return m.P_GENERATOR_MAX_OUTPUT[g] - m.P_GENERATOR_MIN_OUTPUT[g]

    m.P_M_11 = pyo.Param(m.OMEGA_G, rule=m_11_rule, mutable=True)
    m.P_M_12 = pyo.Param(m.OMEGA_G, initialize=1e3, mutable=True)

    def m_21_rule(m, g):
        return m.P_GENERATOR_MAX_OUTPUT[g] - m.P_GENERATOR_MIN_OUTPUT[g]

    m.P_M_21 = pyo.Param(m.OMEGA_G, rule=m_21_rule, mutable=True)
    m.P_M_22 = pyo.Param(m.OMEGA_G, initialize=1e3, mutable=True)

    m.P_M_31 = pyo.Param(m.OMEGA_NM, initialize=float(math.pi), mutable=True)
    m.P_M_32 = pyo.Param(m.OMEGA_NM, initialize=1e3, mutable=True)

    def m_41_rule(m, j):
        if 'REVERSE' in j:
            new_index = j.replace('REVERSE', 'FORWARD')
        elif 'FORWARD' in j:
            new_index = j.replace('FORWARD', 'REVERSE')
        else:
            raise (Exception('REVERSE / FORWARD not in index name'))
        return m.P_NETWORK_INTERCONNECTOR_LIMIT[j] + m.P_NETWORK_INTERCONNECTOR_LIMIT[new_index]

    m.P_M_41 = pyo.Param(m.OMEGA_J, rule=m_41_rule, mutable=True)
    m.P_M_42 = pyo.Param(m.OMEGA_J, initialize=1e3, mutable=True)

    def m_51_rule(m, h):
        return m.P_NETWORK_HVDC_FORWARD_LIMIT[h] - m.P_NETWORK_HVDC_REVERSE_LIMIT[h]

    m.P_M_51 = pyo.Param(m.OMEGA_H, rule=m_51_rule, mutable=True)
    m.P_M_52 = pyo.Param(m.OMEGA_H, initialize=1e3, mutable=True)

    def m_61_rule(m, h):
        return m.P_NETWORK_HVDC_FORWARD_LIMIT[h] - m.P_NETWORK_HVDC_REVERSE_LIMIT[h]

    m.P_M_61 = pyo.Param(m.OMEGA_H, rule=m_61_rule, mutable=True)
    m.P_M_62 = pyo.Param(m.OMEGA_H, initialize=1e3, mutable=True)

    m.P_M_71 = pyo.Param(initialize=1e4, mutable=True)
    m.P_M_72 = pyo.Param(rule=1e4, mutable=True)

    return m

def define_binary_expansion_parameters(m, data):
    """Parameters used when implementing binary expansion discretisation"""

    # Largest integer for baseline selection parameter index
    m.P_BIN_EXP_U = pyo.Param(initialize=data['P_BINARY_EXPANSION_LARGEST_INTEGER'])

    # Minimum emissions intensity baseline
    m.P_BIN_EXP_MIN_BASELINE = pyo.Param(initialize=float(0.8))

    # Maximum emissions intensity baseline
    m.P_BIN_EXP_MAX_BASELINE = pyo.Param(initialize=float(1.3))

    # Emissions intensity baseline increment
    m.P_BIN_EXP_BASELINE_DELTA = pyo.Param(
        initialize=float((m.P_BIN_EXP_MAX_BASELINE - m.P_BIN_EXP_MIN_BASELINE) / (2 ** m.P_BIN_EXP_U)))

    def two_pow_u_rule(m, u):
        """Parameter equal 2^u used when selecting discretised emissions intensity baseline"""
        return float(2 ** u)

    m.P_BIN_EXP_TWO_POW_U = pyo.Param(m.OMEGA_U, rule=two_pow_u_rule)

    # Big-M parameter used used to ensure Z_1=0 when PSI=0, and Z_1=V_DUAL_PERMIT_MARKET when PSI=1
    m.P_BIN_EXP_L_1 = pyo.Param(initialize=1e3)

    def l_2_rule(m, g):
        """ Big-M parameter used to make Z_2=0 when PSI=0, and Z_2=p when PSI=1"""
        return m.P_GENERATOR_MAX_OUTPUT[g] - m.P_GENERATOR_MIN_OUTPUT[g]

    m.P_BIN_EXP_L_2 = pyo.Param(m.OMEGA_G, rule=l_2_rule)

    return m

def define_parameters(m, data):
    """Define parameters"""

    # Policy, network, and generator parameters
    m = define_policy_parameters(m, data)
    m = define_network_parameters(m, data)
    m = define_generator_parameters(m, data)
    m = define_scenario_parameters(m, data)
    m = define_big_m_parameters(m)

    # Binary expansion parameters
    m = define_binary_expansion_parameters(m, data)

    return m

def define_binary_expansion_variables(m):
    """Binary expansion variables"""

    # Binary variables used to determine discretised emissions intensity baseline choice
    m.V_BIN_EXP_PSI = pyo.Var(m.OMEGA_U, within=pyo.Binary)

    # Composite variable - PSI x tau - used to linearise bi-linear term
    m.V_BIN_EXP_Z_1 = pyo.Var(m.OMEGA_U)

    # Composite variable - PSI x p - used to linearise bi-linear term
    m.V_BIN_EXP_Z_2 = pyo.Var(m.OMEGA_S, m.OMEGA_U * m.OMEGA_G)

    return m

def define_primal_variables(m):
    """Define primal variables"""

    # Generator output
    m.V_PRIMAL_GENERATOR_POWER = pyo.Var(m.OMEGA_S, m.OMEGA_G)

    # HVDC link flow
    m.V_PRIMAL_HVDC_FLOW = pyo.Var(m.OMEGA_S, m.OMEGA_H)

    # Node voltage angle
    m.V_PRIMAL_VOLTAGE_ANGLE = pyo.Var(m.OMEGA_S, m.OMEGA_N)

    return m

def define_dual_variables(m):
    """Define dual variables"""

    # Permit market constraint dual variable
    m.V_DUAL_PERMIT_MARKET = pyo.Var()

    # Min power output constraint dual variable
    m.V_DUAL_MU_1 = pyo.Var(m.OMEGA_S, m.OMEGA_G)

    # Max power output constraint dual variable
    m.V_DUAL_MU_2 = pyo.Var(m.OMEGA_S, m.OMEGA_G)

    # Max voltage angle difference constraint dual variable
    m.V_DUAL_MU_3 = pyo.Var(m.OMEGA_S, m.OMEGA_NM)

    # AC link power flow constraint dual variable
    m.V_DUAL_MU_4 = pyo.Var(m.OMEGA_S, m.OMEGA_J)

    # Min HVDC flow constraint dual variable
    m.V_DUAL_MU_5 = pyo.Var(m.OMEGA_S, m.OMEGA_H)

    # Max HVDC flow constraint dual variable
    m.V_DUAL_MU_6 = pyo.Var(m.OMEGA_S, m.OMEGA_H)

    # Reference node voltage angle constraint dual variable
    m.V_DUAL_NU_1 = pyo.Var(m.OMEGA_S, m.OMEGA_N)

    # Node power balance constraint dual variable
    m.V_DUAL_LAMBDA = pyo.Var(m.OMEGA_S, m.OMEGA_N)

    return m

def define_complementarity_constraint_variables(m):
    """Binary variables used in Big-M linearisation"""

    # Permit market constraint binary variable
    m.V_BINARY_PERMIT_MARKET = pyo.Var(within=pyo.Binary)

    # Min power output binary variable
    m.V_BINARY_GAMMA_1 = pyo.Var(m.OMEGA_S, m.OMEGA_G, within=pyo.Binary)

    # Max power output binary variable
    m.V_BINARY_GAMMA_2 = pyo.Var(m.OMEGA_S, m.OMEGA_G, within=pyo.Binary)

    # Max voltage angle difference binary variable
    m.V_BINARY_GAMMA_3 = pyo.Var(m.OMEGA_S, m.OMEGA_NM, within=pyo.Binary)

    # AC link power flow dual variable
    m.V_BINARY_GAMMA_4 = pyo.Var(m.OMEGA_S, m.OMEGA_J, within=pyo.Binary)

    # Min HVDC flow binary variable
    m.V_BINARY_GAMMA_5 = pyo.Var(m.OMEGA_S, m.OMEGA_H, within=pyo.Binary)

    # Max HVDC flow binary variable
    m.V_BINARY_GAMMA_6 = pyo.Var(m.OMEGA_S, m.OMEGA_H, within=pyo.Binary)

    return m

def define_dummy_variables(m):
    """Dummy variables used as a placeholder in objective function"""

    # Dummy variable used in object feasibility mode objective
    m.V_DUMMY = pyo.Var(bounds=(0, 1))

    # Dummy variables used to target given permit price
    m.V_DUMMY_PERMIT_PRICE_TARGET_X_1 = pyo.Var(within=pyo.NonNegativeReals)
    m.V_DUMMY_PERMIT_PRICE_TARGET_X_2 = pyo.Var(within=pyo.NonNegativeReals)

    # Dummy variables used to minimise difference between average price and price target
    m.V_DUMMY_WEIGHTED_RRN_PRICE_X_1 = pyo.Var(within=pyo.NonNegativeReals)
    m.V_DUMMY_WEIGHTED_RRN_PRICE_X_2 = pyo.Var(within=pyo.NonNegativeReals)

    return m

def define_variables(m):
    """Define model variables"""

    # Define primal and dual variables
    m = define_primal_variables(m)
    m = define_dual_variables(m)

    # Variables used in linearised complementarity constraints
    m = define_complementarity_constraint_variables(m)

    # Variables used in binary expansion formulation
    m = define_binary_expansion_variables(m)

    # Dummy variables that may be used as a placeholder within an objective function
    m = define_dummy_variables(m)

    return m

def define_binary_expansion_expressions(m):
    """Binary expansion expressions"""

    # Discretised emissions intensity baseline value
    m.E_BIN_EXP_DISCRETISED_BASELINE = pyo.Expression(
        expr=m.P_BIN_EXP_MIN_BASELINE + (
                m.P_BIN_EXP_BASELINE_DELTA * sum(m.P_BIN_EXP_TWO_POW_U[u] * m.V_BIN_EXP_PSI[u] for u in m.OMEGA_U))
    )

    return m

def define_price_target_expressions(m):
    """Expressions use to compute price targets"""

    # Total revenue from electricity sales
    m.E_TOTAL_REVENUE = pyo.Expression(
        expr=sum(m.P_SCENARIO_DURATION[s] * m.V_DUAL_LAMBDA[s, n] * m.P_NETWORK_DEMAND[s, n]
                 for s in m.OMEGA_S for n in m.OMEGA_N))

    # Total demand
    m.E_TOTAL_DEMAND = pyo.Expression(expr=sum(m.P_SCENARIO_DURATION[s] * m.P_NETWORK_DEMAND[s, n]
                                               for s in m.OMEGA_S for n in m.OMEGA_N))

    # Average price
    m.E_AVERAGE_ELECTRICITY_PRICE = pyo.Expression(expr=m.E_TOTAL_REVENUE / m.E_TOTAL_DEMAND)

    # Weighted RRN price
    m.E_WEIGHTED_RRN_PRICE = pyo.Expression(
        expr=sum(m.P_SCENARIO_DURATION[s] * m.P_NETWORK_REGION_DEMAND_PROPORTION[s, r]
                 * m.V_DUAL_LAMBDA[s, m.P_REGION_RRN_NODE_MAP[r]] for s in m.OMEGA_S for r in m.OMEGA_R))

    return m

def define_expressions(m):
    """Define expressions"""

    # Binary expansion expressions
    m = define_binary_expansion_expressions(m)

    # Expressions used in price targeting objective functions
    m = define_price_target_expressions(m)

    return m

def define_binary_expansion_constraints(m):
    """Constraints used in binary expansion formulation"""

    # Constraints ensures Z_1=0 when PSI=0 and Z_1=permit_price when PSI=1
    def z_1_constraint_1_rule(m, u):
        return 0 <= m.V_DUAL_PERMIT_MARKET - m.V_BIN_EXP_Z_1[u]

    m.C_Z_1_CONSTRAINT_1 = pyo.Constraint(m.OMEGA_U, rule=z_1_constraint_1_rule)

    def z_1_constraint_2_rule(m, u):
        return m.V_DUAL_PERMIT_MARKET - m.V_BIN_EXP_Z_1[u] <= m.P_BIN_EXP_L_1 * (1 - m.V_BIN_EXP_PSI[u])

    m.C_Z_1_CONSTRAINT_2 = pyo.Constraint(m.OMEGA_U, rule=z_1_constraint_2_rule)

    def z_1_constraint_3_rule(m, u):
        return 0 <= m.V_BIN_EXP_Z_1[u]

    m.C_Z_1_CONSTRAINT_3 = pyo.Constraint(m.OMEGA_U, rule=z_1_constraint_3_rule)

    def z_1_constraint_4_rule(m, u):
        return m.V_BIN_EXP_Z_1[u] <= m.P_BIN_EXP_L_1 * m.V_BIN_EXP_PSI[u]

    m.C_Z_1_CONSTRAINT_4 = pyo.Constraint(m.OMEGA_U, rule=z_1_constraint_4_rule)

    # Constraints ensuring z_2=0 when PSI=0, and z_2=p when PSI=1
    def z_2_constraint_1_rule(m, s, u, g):
        return 0 <= m.V_PRIMAL_GENERATOR_POWER[s, g] - m.V_BIN_EXP_Z_2[s, u, g]

    m.C_Z_2_CONSTRAINT_1 = pyo.Constraint(m.OMEGA_S, m.OMEGA_U * m.OMEGA_G, rule=z_2_constraint_1_rule)

    def z_2_constraint_2_rule(m, s, u, g):
        return m.V_PRIMAL_GENERATOR_POWER[s, g] - m.V_BIN_EXP_Z_2[s, u, g] <= m.P_BIN_EXP_L_2[g] * (
                1 - m.V_BIN_EXP_PSI[u])

    m.C_Z_2_CONSTRAINT_2 = pyo.Constraint(m.OMEGA_S, m.OMEGA_U * m.OMEGA_G, rule=z_2_constraint_2_rule)

    def z_2_constraint_3_rule(m, s, u, g):
        return 0 <= m.V_BIN_EXP_Z_2[s, u, g]

    m.C_Z_2_CONSTRAINT_3 = pyo.Constraint(m.OMEGA_S, m.OMEGA_U * m.OMEGA_G, rule=z_2_constraint_3_rule)

    def z_2_constraint_4_rule(m, s, u, g):
        return m.V_BIN_EXP_Z_2[s, u, g] <= m.P_BIN_EXP_L_2[g] * m.V_BIN_EXP_PSI[u]

    m.C_Z_2_CONSTRAINT_4 = pyo.Constraint(m.OMEGA_S, m.OMEGA_U * m.OMEGA_G, rule=z_2_constraint_4_rule)

    return m

def define_first_order_condition_constraints(m):
    """Define constraints constituting the model's first order conditions"""

    def foc_1_rule(m, s, g):
        """To be activated if baseline is fixed"""
        return (m.P_GENERATOR_SRMC[g]
                + ((m.P_GENERATOR_EMISSIONS_INTENSITY[g] - m.P_POLICY_FIXED_BASELINE) * m.V_DUAL_PERMIT_MARKET)
                - m.V_DUAL_MU_1[s, g]
                + m.V_DUAL_MU_2[s, g]
                - m.V_DUAL_LAMBDA[s, m.P_GENERATOR_NODE[g]] == 0)

    m.C_FOC_1 = pyo.Constraint(m.OMEGA_S, m.OMEGA_G, rule=foc_1_rule)

    def foc_1_linearised_rule(m, s, g):
        """To be activated if baseline is discretised"""
        return (m.P_GENERATOR_SRMC[g]
                + ((m.P_GENERATOR_EMISSIONS_INTENSITY[g] - m.P_BIN_EXP_MIN_BASELINE) * m.V_DUAL_PERMIT_MARKET)
                - (m.P_BIN_EXP_BASELINE_DELTA * sum(m.P_BIN_EXP_TWO_POW_U[u] * m.V_BIN_EXP_Z_1[u] for u in m.OMEGA_U))
                - m.V_DUAL_MU_1[s, g]
                + m.V_DUAL_MU_2[s, g]
                - m.V_DUAL_LAMBDA[s, m.P_GENERATOR_NODE[g]]
                == 0)

    m.C_FOC_1_LINEARISED = pyo.Constraint(m.OMEGA_S, m.OMEGA_G, rule=foc_1_linearised_rule)

    def foc_2_rule(m, s, n):
        """Constraint relating to voltage angles"""

        term_1 = sum(m.V_DUAL_MU_3[s, n, j]
                     - m.V_DUAL_MU_3[s, j, n]
                     + (m.V_DUAL_LAMBDA[s, n] * m.P_NETWORK_BRANCH_SUSCEPTANCE[n, j])
                     - (m.V_DUAL_LAMBDA[s, j] * m.P_NETWORK_BRANCH_SUSCEPTANCE[j, n]) for j in m.P_NETWORK_GRAPH[n])

        term_2 = m.V_DUAL_NU_1[s, n] * m.P_NETWORK_REFERENCE_NODE_INDICATOR[n]

        term_3 = sum(
            m.P_NETWORK_BRANCH_SUSCEPTANCE[m.P_NETWORK_INTERCONNECTOR_BRANCH_NODES_MAP[l]]
            * m.V_DUAL_MU_4[s, j]
            * m.P_NETWORK_INTERCONNECTOR_INCIDENCE_MAT[n, l]
            for j in m.OMEGA_J for l in m.P_NETWORK_INTERCONNECTOR_BRANCH_ID_MAP[j])

        return term_1 + term_2 + term_3 == 0

    m.C_FOC_2 = pyo.Constraint(m.OMEGA_S, m.OMEGA_N, rule=foc_2_rule)

    def foc_3_rule(m, s, h):
        """Constraint relating to HVDC flows"""

        # From and to nodes for given link, h
        from_node = m.P_NETWORK_HVDC_FROM_NODE[h]
        to_node = m.P_NETWORK_HVDC_TO_NODE[h]

        return ((m.P_NETWORK_HVDC_INCIDENCE_MAT[from_node, h] * m.V_DUAL_LAMBDA[s, from_node])
                + (m.P_NETWORK_HVDC_INCIDENCE_MAT[to_node, h] * m.V_DUAL_LAMBDA[s, to_node])
                - m.V_DUAL_MU_5[s, h]
                + m.V_DUAL_MU_6[s, h] == 0)

    m.C_FOC_3 = pyo.Constraint(m.OMEGA_S, m.OMEGA_H, rule=foc_3_rule)

    return m

def define_network_constraint(m):
    """Define network constraints"""

    def reference_node_voltage_angle_rule(m, s, n):
        """Fix reference node voltage angles to 0"""

        if m.P_NETWORK_REFERENCE_NODE_INDICATOR[n] == 1:
            return m.V_PRIMAL_VOLTAGE_ANGLE[s, n] == 0
        else:
            return pyo.Constraint.Skip

    m.C_REFERENCE_NODE_VOLTAGE_ANGLE = pyo.Constraint(m.OMEGA_S, m.OMEGA_N, rule=reference_node_voltage_angle_rule)

    def power_balance_rule(m, s, n):
        """Power balance at each node"""

        # Total injection from generators connected to node
        generator_injection = sum(m.V_PRIMAL_GENERATOR_POWER[s, g] for g in m.P_NETWORK_NODE_GENERATOR_MAP[n])

        # Injection from AC branches
        ac_injection = sum(
            m.P_NETWORK_BRANCH_SUSCEPTANCE[n, j] * (m.V_PRIMAL_VOLTAGE_ANGLE[s, n] - m.V_PRIMAL_VOLTAGE_ANGLE[s, j])
            for j in m.P_NETWORK_GRAPH[n])

        # Injection from HVDC branches
        hvdc_injection = sum(m.P_NETWORK_HVDC_INCIDENCE_MAT[n, h] * m.V_PRIMAL_HVDC_FLOW[s, h] for h in m.OMEGA_H)

        return (m.P_NETWORK_DEMAND[s, n]
                - m.P_NETWORK_FIXED_INJECTION[s, n]
                - generator_injection
                + ac_injection
                + hvdc_injection
                == 0)

    m.C_POWER_BALANCE = pyo.Constraint(m.OMEGA_S, m.OMEGA_N, rule=power_balance_rule)

    return m

def define_linearised_complementarity_constraints(m):
    """Linearised complementarity constraints"""

    def lin_comp_1_1_rule(m, s, g):
        """Min power output constraint 1"""

        return m.P_GENERATOR_MIN_OUTPUT[g] - m.V_PRIMAL_GENERATOR_POWER[s, g] <= 0

    m.C_LIN_COMP_1_1 = pyo.Constraint(m.OMEGA_S, m.OMEGA_G, rule=lin_comp_1_1_rule)

    def lin_comp_1_2_rule(m, s, g):
        """Min power output constraint 2"""

        return m.V_DUAL_MU_1[s, g] >= 0

    m.C_LIN_COMP_1_2 = pyo.Constraint(m.OMEGA_S, m.OMEGA_G, rule=lin_comp_1_2_rule)

    def lin_comp_1_3_rule(m, s, g):
        """Min power output constraint 3"""

        return m.V_PRIMAL_GENERATOR_POWER[s, g] - m.P_GENERATOR_MIN_OUTPUT[g] <= m.V_BINARY_GAMMA_1[s, g] * m.P_M_11[g]

    m.C_LIN_COMP_1_3 = pyo.Constraint(m.OMEGA_S, m.OMEGA_G, rule=lin_comp_1_3_rule)

    def lin_comp_1_4_rule(m, s, g):
        """Min power output constraint 4"""

        return m.V_DUAL_MU_1[s, g] <= (1 - m.V_BINARY_GAMMA_1[s, g]) * m.P_M_12[g]

    m.C_LIN_COMP_1_4 = pyo.Constraint(m.OMEGA_S, m.OMEGA_G, rule=lin_comp_1_4_rule)

    def lin_comp_2_1_rule(m, s, g):
        """Max power output constraint 1"""

        return m.V_PRIMAL_GENERATOR_POWER[s, g] - m.P_GENERATOR_MAX_OUTPUT[g] <= 0

    m.C_LIN_COMP_2_1 = pyo.Constraint(m.OMEGA_S, m.OMEGA_G, rule=lin_comp_2_1_rule)

    def lin_comp_2_2_rule(m, s, g):
        """Max power output constraint 2"""

        return m.V_DUAL_MU_2[s, g] >= 0

    m.C_LIN_COMP_2_2 = pyo.Constraint(m.OMEGA_S, m.OMEGA_G, rule=lin_comp_2_2_rule)

    def lin_comp_2_3_rule(m, s, g):
        """Max power output constraint 3"""

        return (m.P_GENERATOR_MAX_OUTPUT[g] - m.V_PRIMAL_GENERATOR_POWER[s, g]
                <= m.V_BINARY_GAMMA_2[s, g] * m.P_M_21[g])

    m.C_LIN_COMP_2_3 = pyo.Constraint(m.OMEGA_S, m.OMEGA_G, rule=lin_comp_2_3_rule)

    def lin_comp_2_4_rule(m, s, g):
        """Max power output constraint 4"""

        return m.V_DUAL_MU_2[s, g] <= (1 - m.V_BINARY_GAMMA_2[s, g]) * m.P_M_22[g]

    m.C_LIN_COMP_2_4 = pyo.Constraint(m.OMEGA_S, m.OMEGA_G, rule=lin_comp_2_4_rule)

    def lin_comp_3_1_rule(m, s, n, j):
        """Max voltage angle difference between connected nodes constraint 1"""

        return (m.V_PRIMAL_VOLTAGE_ANGLE[s, n] - m.V_PRIMAL_VOLTAGE_ANGLE[s, j]
                - m.P_NETWORK_VOLTAGE_ANGLE_MAX_DIFFERENCE
                <= 0)

    m.C_LIN_COMP_3_1 = pyo.Constraint(m.OMEGA_S, m.OMEGA_NM, rule=lin_comp_3_1_rule)

    def lin_comp_3_2_rule(m, s, n, j):
        """Max voltage angle difference between connected nodes constraint 2"""

        return m.V_DUAL_MU_3[s, n, j] >= 0

    m.C_LIN_COMP_3_2 = pyo.Constraint(m.OMEGA_S, m.OMEGA_NM, rule=lin_comp_3_2_rule)

    def lin_comp_3_3_rule(m, s, n, j):
        """Max voltage angle difference between connected nodes constraint 3"""

        return (m.P_NETWORK_VOLTAGE_ANGLE_MAX_DIFFERENCE + m.V_PRIMAL_VOLTAGE_ANGLE[s, j]
                - m.V_PRIMAL_VOLTAGE_ANGLE[s, n] <= m.V_BINARY_GAMMA_3[s, n, j] * m.P_M_31[n, j])

    m.C_LIN_COMP_3_3 = pyo.Constraint(m.OMEGA_S, m.OMEGA_NM, rule=lin_comp_3_3_rule)

    def lin_comp_3_4_rule(m, s, n, j):
        """Max voltage angle difference between connected nodes constraint 4"""

        return m.V_DUAL_MU_3[s, n, j] <= (1 - m.V_BINARY_GAMMA_3[s, n, j]) * m.P_M_32[n, j]

    m.C_LIN_COMP_3_4 = pyo.Constraint(m.OMEGA_S, m.OMEGA_NM, rule=lin_comp_3_4_rule)

    def lin_comp_4_1_rule(m, s, j):
        """Interconnector flow limits constraint 1"""

        # Links constituting interconnector, j
        branches = m.P_NETWORK_INTERCONNECTOR_BRANCH_ID_MAP[j]
        branch_links = [m.P_NETWORK_INTERCONNECTOR_BRANCH_NODES_MAP[i] for i in branches]

        return (sum(m.P_NETWORK_BRANCH_SUSCEPTANCE[n, p]
                    * (m.V_PRIMAL_VOLTAGE_ANGLE[s, n] - m.V_PRIMAL_VOLTAGE_ANGLE[s, p])
                    for n, p in branch_links) - m.P_NETWORK_INTERCONNECTOR_LIMIT[j]
                <= 0)

    m.C_LIN_COMP_4_1 = pyo.Constraint(m.OMEGA_S, m.OMEGA_J, rule=lin_comp_4_1_rule)

    def lin_comp_4_2_rule(m, s, j):
        """Interconnector flow limits constraint 2"""

        return m.V_DUAL_MU_4[s, j] >= 0

    m.C_LIN_COMP_4_2 = pyo.Constraint(m.OMEGA_S, m.OMEGA_J, rule=lin_comp_4_2_rule)

    def lin_comp_4_3_rule(m, s, j):
        """Interconnector flow limits constraint 3"""

        # Links constituting interconnector, j
        branches = m.P_NETWORK_INTERCONNECTOR_BRANCH_ID_MAP[j]
        branch_links = [m.P_NETWORK_INTERCONNECTOR_BRANCH_NODES_MAP[i] for i in branches]

        return (m.P_NETWORK_INTERCONNECTOR_LIMIT[j]
                - sum(m.P_NETWORK_BRANCH_SUSCEPTANCE[n, p]
                      * (m.V_PRIMAL_VOLTAGE_ANGLE[s, n] - m.V_PRIMAL_VOLTAGE_ANGLE[s, p]) for n, p in branch_links)
                <= m.V_BINARY_GAMMA_4[s, j] * m.P_M_41[j])

    m.C_LIN_COMP_4_3 = pyo.Constraint(m.OMEGA_S, m.OMEGA_J, rule=lin_comp_4_3_rule)

    def lin_comp_4_4_rule(m, s, j):
        """Interconnector flow limits constraint 4"""

        return m.V_DUAL_MU_4[s, j] <= (1 - m.V_BINARY_GAMMA_4[s, j]) * m.P_M_42[j]

    m.C_LIN_COMP_4_4 = pyo.Constraint(m.OMEGA_S, m.OMEGA_J, rule=lin_comp_4_4_rule)

    def lin_comp_5_1_rule(m, s, h):
        """HVDC reverse flow limits constraint 1"""

        return m.P_NETWORK_HVDC_REVERSE_LIMIT[h] - m.V_PRIMAL_HVDC_FLOW[s, h] <= 0

    m.C_LIN_COMP_5_1 = pyo.Constraint(m.OMEGA_S, m.OMEGA_H, rule=lin_comp_5_1_rule)

    def lin_comp_5_2_rule(m, s, h):
        """HVDC reverse flow limits constraint 2"""

        return m.V_DUAL_MU_5[s, h] >= 0

    m.C_LIN_COMP_5_2 = pyo.Constraint(m.OMEGA_S, m.OMEGA_H, rule=lin_comp_5_2_rule)

    def lin_comp_5_3_rule(m, s, h):
        """HVDC reverse flow limits constraint 3"""

        return m.V_PRIMAL_HVDC_FLOW[s, h] - m.P_NETWORK_HVDC_REVERSE_LIMIT[h] <= m.V_BINARY_GAMMA_5[s, h] * m.P_M_51[h]

    m.C_LIN_COMP_5_3 = pyo.Constraint(m.OMEGA_S, m.OMEGA_H, rule=lin_comp_5_3_rule)

    def lin_comp_5_4_rule(m, s, h):
        """HVDC reverse flow limits constraint 4"""

        return m.V_DUAL_MU_5[s, h] <= (1 - m.V_BINARY_GAMMA_5[s, h]) * m.P_M_52[h]

    m.C_LIN_COMP_5_4 = pyo.Constraint(m.OMEGA_S, m.OMEGA_H, rule=lin_comp_5_4_rule)

    def lin_comp_6_1_rule(m, s, h):
        """HVDC forward flow limits constraint 1"""

        return m.V_PRIMAL_HVDC_FLOW[s, h] - m.P_NETWORK_HVDC_FORWARD_LIMIT[h] <= 0

    m.C_LIN_COMP_6_1 = pyo.Constraint(m.OMEGA_S, m.OMEGA_H, rule=lin_comp_6_1_rule)

    def lin_comp_6_2_rule(m, s, h):
        """HVDC forward flow limits constraint 2"""

        return m.V_DUAL_MU_6[s, h] >= 0

    m.C_LIN_COMP_6_2 = pyo.Constraint(m.OMEGA_S, m.OMEGA_H, rule=lin_comp_6_2_rule)

    def lin_comp_6_3_rule(m, s, h):
        """HVDC forward flow limits constraint 3"""

        return m.P_NETWORK_HVDC_FORWARD_LIMIT[h] - m.V_PRIMAL_HVDC_FLOW[s, h] <= m.V_BINARY_GAMMA_6[s, h] * m.P_M_61[h]

    m.C_LIN_COMP_6_3 = pyo.Constraint(m.OMEGA_S, m.OMEGA_H, rule=lin_comp_6_3_rule)

    def lin_comp_6_4_rule(m, s, h):
        """HVDC forward flow limits constraint 4"""

        return m.V_DUAL_MU_6[s, h] <= (1 - m.V_BINARY_GAMMA_6[s, h]) * m.P_M_62[h]

    m.C_LIN_COMP_6_4 = pyo.Constraint(m.OMEGA_S, m.OMEGA_H, rule=lin_comp_6_4_rule)

    return m

def define_permit_market_constraints(m):
    """Define permit market constraints"""

    def permit_market_1_rule(m):
        """Permit market complementarity constraint 1 - use when baseline is fixed"""

        return (sum(m.P_SCENARIO_DURATION[s] * ((m.P_GENERATOR_EMISSIONS_INTENSITY[g] - m.P_POLICY_FIXED_BASELINE)
                                                * m.V_PRIMAL_GENERATOR_POWER[s, g])
                    for g in m.OMEGA_G for s in m.OMEGA_S)
                <= 0)

    m.C_PERMIT_MARKET_1 = pyo.Constraint(rule=permit_market_1_rule)

    def permit_market_1_linearised_rule(m):
        """Permit market complementarity constraint 1 - use when baseline is variable"""

        return (sum(m.P_SCENARIO_DURATION[s]
                    * (((m.P_GENERATOR_EMISSIONS_INTENSITY[g] - m.P_BIN_EXP_MIN_BASELINE)
                        * m.V_PRIMAL_GENERATOR_POWER[s, g])
                       - (m.P_BIN_EXP_BASELINE_DELTA * sum(m.P_BIN_EXP_TWO_POW_U[u] * m.V_BIN_EXP_Z_2[s, u, g]
                                                           for u in m.OMEGA_U)))
                    for g in m.OMEGA_G for s in m.OMEGA_S)
                <= 0)

    m.C_PERMIT_MARKET_1_LINEARISED = pyo.Constraint(rule=permit_market_1_linearised_rule)

    def permit_market_2_rule(m):
        """Permit market complementarity constraint 2"""

        return m.V_DUAL_PERMIT_MARKET >= 0

    m.C_PERMIT_MARKET_2 = pyo.Constraint(rule=permit_market_2_rule)

    def permit_market_3_rule(m):
        """Permit market complementarity constraint 3"""

        return sum(m.P_SCENARIO_DURATION[s]
                   * ((m.P_POLICY_FIXED_BASELINE - m.P_GENERATOR_EMISSIONS_INTENSITY[g])
                      * m.V_PRIMAL_GENERATOR_POWER[s, g])
                   for g in m.OMEGA_G for s in m.OMEGA_S) <= (m.V_BINARY_PERMIT_MARKET * m.P_M_71)

    m.C_PERMIT_MARKET_3 = pyo.Constraint(rule=permit_market_3_rule)

    def permit_market_3_linearised_rule(m):
        """Permit market complementarity constraint 3 - linearised"""

        return (sum(m.P_SCENARIO_DURATION[s] * (((m.P_BIN_EXP_MIN_BASELINE - m.P_GENERATOR_EMISSIONS_INTENSITY[g])
                                                 * m.V_PRIMAL_GENERATOR_POWER[s, g])
                                                + (m.P_BIN_EXP_BASELINE_DELTA
                                                   * sum(m.P_BIN_EXP_TWO_POW_U[u] * m.V_BIN_EXP_Z_2[s, u, g]
                                                         for u in m.OMEGA_U))) for g in m.OMEGA_G for s in m.OMEGA_S)
                <= m.V_BINARY_PERMIT_MARKET * m.P_M_71)

    m.C_PERMIT_MARKET_3_LINEARISED = pyo.Constraint(rule=permit_market_3_linearised_rule)

    def permit_market_4_rule(m):
        """Permit market complementarity constraint 4"""

        return m.V_DUAL_PERMIT_MARKET <= (1 - m.V_BINARY_PERMIT_MARKET) * m.P_M_72

    m.C_PERMIT_MARKET_4 = pyo.Constraint(rule=permit_market_4_rule)

    return m

def define_permit_price_targeting_constraints(m):
    """Constraints used to get the absolute difference between the permit price and some target"""

    # Constraints to minimise difference between permit price and target
    m.C_PERMIT_PRICE_TARGET_CONSTRAINT_1 = pyo.Constraint(
        expr=m.V_DUMMY_PERMIT_PRICE_TARGET_X_1 >= m.P_POLICY_PERMIT_PRICE_TARGET - m.V_DUAL_PERMIT_MARKET)

    m.C_PERMIT_PRICE_TARGET_CONSTRAINT_2 = pyo.Constraint(
        expr=m.V_DUMMY_PERMIT_PRICE_TARGET_X_2 >= m.V_DUAL_PERMIT_MARKET - m.P_POLICY_PERMIT_PRICE_TARGET)

    return m

def define_weighted_rrn_price_targeting_constraints(m):
    """Constraints used to get absolute difference between weighted RRN price and some target"""

    # Constraints used to minimise difference between RRN price target and RRN price
    m.C_WEIGHTED_RRN_PRICE_TARGET_CONSTRAINT_1 = pyo.Constraint(
        expr=m.V_DUMMY_WEIGHTED_RRN_PRICE_X_1 >= m.E_WEIGHTED_RRN_PRICE - m.P_POLICY_WEIGHTED_RRN_PRICE_TARGET)

    m.C_WEIGHTED_RRN_PRICE_TARGET_CONSTRAINT_2 = pyo.Constraint(
        expr=m.V_DUMMY_WEIGHTED_RRN_PRICE_X_2 >= m.P_POLICY_WEIGHTED_RRN_PRICE_TARGET - m.E_WEIGHTED_RRN_PRICE)

    return m

def define_constraints(m):
    """Define constraints"""

    m = define_binary_expansion_constraints(m)
    m = define_first_order_condition_constraints(m)
    m = define_network_constraint(m)
    m = define_linearised_complementarity_constraints(m)
    m = define_permit_market_constraints(m)

    # Constraints used to obtain absolute values between variables and some target
    m = define_permit_price_targeting_constraints(m)
    m = define_weighted_rrn_price_targeting_constraints(m)

    return m

def define_objective_functions(m):
    """Define objective functions"""

    # Use if trying to find a feasible solution for a fixed baseline
    m.O_FEASIBILITY = pyo.Objective(expr=m.V_DUMMY, sense=pyo.minimize)

    # Objective function
    m.O_PERMIT_PRICE_TARGET = pyo.Objective(expr=m.V_DUMMY_PERMIT_PRICE_TARGET_X_1 + m.V_DUMMY_PERMIT_PRICE_TARGET_X_2)

    # Weighted RRN price targeting objective function
    m.O_WEIGHTED_RRN_PRICE_TARGET = pyo.Objective(
        expr=m.V_DUMMY_WEIGHTED_RRN_PRICE_X_1 + m.V_DUMMY_WEIGHTED_RRN_PRICE_X_2)

    return m

def apply_pu_scaling(m):
    """Attempt to improve numerical stability by scaling selected parameters"""

    # Scale branch susceptance
    for i in m.P_NETWORK_BRANCH_SUSCEPTANCE.keys():
        m.P_NETWORK_BRANCH_SUSCEPTANCE[i] = m.P_NETWORK_BRANCH_SUSCEPTANCE[i] / m.P_NETWORK_BASE_POWER

    # HVDC limits
    for i in m.P_NETWORK_HVDC_FORWARD_LIMIT.keys():
        m.P_NETWORK_HVDC_FORWARD_LIMIT[i] = m.P_NETWORK_HVDC_FORWARD_LIMIT[i] / m.P_NETWORK_BASE_POWER

    for i in m.P_NETWORK_HVDC_REVERSE_LIMIT.keys():
        m.P_NETWORK_HVDC_REVERSE_LIMIT[i] = m.P_NETWORK_HVDC_REVERSE_LIMIT[i] / m.P_NETWORK_BASE_POWER

    # Generator max output
    for i in m.P_GENERATOR_MAX_OUTPUT.keys():
        m.P_GENERATOR_MAX_OUTPUT[i] = m.P_GENERATOR_MAX_OUTPUT[i] / m.P_NETWORK_BASE_POWER

    # Generator short run marginal costs
    for i in m.P_GENERATOR_SRMC.keys():
        m.P_GENERATOR_SRMC[i] = m.P_GENERATOR_SRMC[i] / m.P_NETWORK_BASE_POWER

    # Interconnector flow limits
    for i in m.P_NETWORK_INTERCONNECTOR_LIMIT.keys():
        m.P_NETWORK_INTERCONNECTOR_LIMIT[i] = m.P_NETWORK_INTERCONNECTOR_LIMIT[i] / m.P_NETWORK_BASE_POWER

    # Fixed injections
    for i in m.P_NETWORK_FIXED_INJECTION.keys():
        m.P_NETWORK_FIXED_INJECTION[i] = m.P_NETWORK_FIXED_INJECTION[i] / m.P_NETWORK_BASE_POWER

    # Demand
    for i in m.P_NETWORK_DEMAND.keys():
        m.P_NETWORK_DEMAND[i] = m.P_NETWORK_DEMAND[i] / m.P_NETWORK_BASE_POWER

    # Complementarity constraints
    big_m_parameters = [f'P_M_{i}{j}' for i in range(1, 8) for j in range(1, 3)]
    ignore_scaling = ['P_M_31']
    for i in big_m_parameters:
        if i in ignore_scaling:
            continue
        else:
            for j in m.__getattribute__(i).keys():
                m.__getattribute__(i)[j] = m.__getattribute__(i)[j] / m.P_NETWORK_BASE_POWER

    # Scale price targets
    m.P_POLICY_PERMIT_PRICE_TARGET = m.P_POLICY_PERMIT_PRICE_TARGET.value / m.P_NETWORK_BASE_POWER.value
    m.P_POLICY_WEIGHTED_RRN_PRICE_TARGET = m.P_POLICY_WEIGHTED_RRN_PRICE_TARGET.value / m.P_NETWORK_BASE_POWER.value

    return m

def construct_model(data, use_pu):
    """Construct model"""

    # Initialise model
    m = pyo.ConcreteModel()

    # Define sets and parameters
    m = define_sets(m, data)
    m = define_parameters(m, data)

    # Optionally apply per-unit scaling - attempt to improve numerical stability
    if use_pu:
        m = apply_pu_scaling(m)

    # Define variables, expressions, and constraints
    m = define_variables(m)
    m = define_expressions(m)
    m = define_constraints(m)

    # Define objective functions
    m = define_objective_functions(m)

    return m

def configure_feasibility_model(m):
    """Configure model to accept a fixed emissions intensity baseline"""

    # Deactivate price targeting objective functions
    m.O_PERMIT_PRICE_TARGET.deactivate()
    m.O_WEIGHTED_RRN_PRICE_TARGET.deactivate()

    # Deactivate linearised constraints - those with discretised baseline
    m.C_FOC_1_LINEARISED.deactivate()
    m.C_PERMIT_MARKET_1_LINEARISED.deactivate()
    m.C_PERMIT_MARKET_3_LINEARISED.deactivate()

    # Deactivate constraints used for binary expansion linearisation
    m.C_Z_1_CONSTRAINT_1.deactivate()
    m.C_Z_1_CONSTRAINT_2.deactivate()
    m.C_Z_1_CONSTRAINT_3.deactivate()
    m.C_Z_1_CONSTRAINT_4.deactivate()
    m.C_Z_2_CONSTRAINT_1.deactivate()
    m.C_Z_2_CONSTRAINT_2.deactivate()
    m.C_Z_2_CONSTRAINT_3.deactivate()
    m.C_Z_2_CONSTRAINT_4.deactivate()

    return m

def configure_weighted_rrn_price_targeting_model(m):
    """Configure mode to target weighted RRN prices"""

    # Ensure weighted RRN price targeting objective is the only one active
    m.O_FEASIBILITY.deactivate()
    m.O_PERMIT_PRICE_TARGET.deactivate()

    # Deactivate first order condition and permit market constraints that have not been linearised
    m.C_FOC_1.deactivate()
    m.C_PERMIT_MARKET_1.deactivate()
    m.C_PERMIT_MARKET_3.deactivate()

    return m

def configure_permit_price_targeting_model(m):
    """Configure permit price targeting model"""

    # Ensure permit price targeting objective is the only one active
    m.O_FEASIBILITY.deactivate()
    m.O_WEIGHTED_RRN_PRICE_TARGET.deactivate()

    # Deactivate first order condition and permit market constraints that have not been linearised
    m.C_FOC_1.deactivate()
    m.C_PERMIT_MARKET_1.deactivate()
    m.C_PERMIT_MARKET_3.deactivate()

    return m

def configure_model(m, options):
    """Activate / deactivate constraints and objectives depending on the case to be analysed"""

    if options['mode'] == 'feasibility':
        m = configure_feasibility_model(m)

    elif options['mode'] == 'weighted_rrn_price_target':
        m = configure_weighted_rrn_price_targeting_model(m)

    elif options['mode'] == 'permit_price_target':
        m = configure_permit_price_targeting_model(m)

    else:
        raise Exception(f"Unexpected mode: {options['mode']}")

    return m

def solve_model(m):
    """Solve model"""

    opt = pyo.SolverFactory('cplex', solver_io='mps')

    # Solver options
    options = {
        'mip tolerances absmipgap': 1e-6,
        'emphasis mip': 1,  # Emphasise feasibility
        # 'timelimit': 5000
    }

    # Solve model
    opt.solve(m, keepfiles=False, tee=True, options=options)

    return m

def extract_value(v):
    """Extract model value"""

    if type(v) in [int, float]:
        return v
    else:
        return v.value

def get_solution(m, params):
    """Return model solution"""

    # Extract selected model parameters - used to check case file was correctly specified
    model_parameters = {
        'P_POLICY_PERMIT_PRICE_TARGET': m.P_POLICY_PERMIT_PRICE_TARGET.value,
        'P_POLICY_WEIGHTED_RRN_PRICE_TARGET': m.P_POLICY_PERMIT_PRICE_TARGET.value,
        'P_SCENARIO_DURATION': {k: extract_value(v) for k, v in m.P_SCENARIO_DURATION.items()},
        'P_GENERATOR_SRMC': {k: extract_value(v) for k, v in m.P_GENERATOR_SRMC.items()},
        'P_GENERATOR_EMISSIONS_INTENSITY': {k: extract_value(v) for k, v in m.P_GENERATOR_EMISSIONS_INTENSITY.items()},
    }

    # Extract selected variables and expressions
    try:
        variables = {
            'V_PRIMAL_GENERATOR_POWER': {k: extract_value(v) for k, v in m.V_PRIMAL_GENERATOR_POWER.items()},
            'V_DUAL_PERMIT_MARKET': m.V_DUAL_PERMIT_MARKET.value,
        }
    except ValueError:
        variables = {}

    try:
        expressions = {
            'E_WEIGHTED_RRN_PRICE': m.E_WEIGHTED_RRN_PRICE.expr(),
            'E_AVERAGE_ELECTRICITY_PRICE': m.E_AVERAGE_ELECTRICITY_PRICE.expr(),
        }
    except ValueError:
        expressions = {}

    # Extract discretised baseline - can only be extracted when targeting prices
    try:
        price_targeting_expressions = {
            'E_BIN_EXP_DISCRETISED_BASELINE': m.E_BIN_EXP_DISCRETISED_BASELINE.expr(),
        }
    except ValueError:
        price_targeting_expressions = {}

    # Combine into single dictionary
    solution = {**model_parameters, **variables, **expressions, **price_targeting_expressions}

    # Combine into single dictionary
    output = {
        'options': params,
        'solution': solution
    }

    return output

if __name__ == '__main__':
    # Data and output directories
    data_directory = os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir, 'data')
    scenario_directory = os.path.join(os.path.dirname(__file__), os.path.pardir, '1_create_scenarios', 'output')
    tmp_directory = os.path.join(os.path.dirname(__file__), 'utils', 'tmp')

    model_options = {
        'parameters': {
            'P_BINARY_EXPANSION_LARGEST_INTEGER': 10,
            'P_POLICY_FIXED_BASELINE': 1.2,
            'P_POLICY_PERMIT_PRICE_TARGET': 30,
            'P_POLICY_WEIGHTED_RRN_PRICE_TARGET': 20,
        },
        'mode': 'feasibility',
    }

    # Get case data
    case_data = get_case_data(data_directory, scenario_directory, tmp_directory, model_options, use_cache=True)

    # Construct model
    model = construct_model(case_data, use_pu=True)
    model = configure_model(model, model_options)
    model = solve_model(model)


import time
from functools import wraps
import datetime
import decimal
import hashlib
from urllib.parse import urlparse, parse_qsl

from bencode import bencode
from flask import request, Response
from flask.json import JSONEncoder, JSONDecoder
# Define custom JSONEncoder for the ISO Datetime format
from flask_jwt_extended import get_jwt_identity, get_jwt_claims
from flask_restful.reqparse import Namespace
from json.decoder import WHITESPACE

settings = {
    'time_until_inactive': 30 * 60,  # 30 minutes
    'seed_interval': 300,  # 5 minutes
    'peer_interval': 30,  # 30 seconds
    'numwant_default': 30,  # 30 should be plent as stated by spec
    'compact_default': 0,  # 1 for bandwidth, 0 for compatability
    'no_peer_id_default': 0,  # 1 for bandwidth, 0 for compatability

    # Don't change these. I mean you can but its pointless
    'seed': 'S',
    'peer': 'P',
}

class JSONEncoder(JSONEncoder):
    def default(self, obj):
        try:
            if isinstance(obj, datetime.date):
                return obj.isoformat()
            elif isinstance(obj, decimal.Decimal):
                return float(obj)
            iterable = iter(obj)
        except TypeError:
            pass
        else:
            return list(iterable)
        return JSONEncoder.default(self, obj)

class JSONDecoder(JSONDecoder):
    unicode_replacements = {
        '\u2018': "'", '\u2019': "'"
    }

    def __init__(self, *args, **kwargs):
        self.orig_obj_hook = kwargs.pop("object_hook", None)
        super(JSONDecoder, self).__init__(*args, object_hook=self.custom_obj_hook, strict=False, **kwargs)

    def decode(self, s, _w=WHITESPACE.match):
        for rk in self.unicode_replacements.keys():
            s = s.replace(rk, self.unicode_replacements.get(rk))
        # if max(s) > u'\u00FF':
        #     print(f'Unicode out of range in {s.index(max(s))}. Deleting that character and continuing')
        return super().decode(s, _w)

    def custom_obj_hook(self, dct):
        # Calling custom decode function:4
        if self.orig_obj_hook:  # Do we have another hook to call?
            return self.orig_obj_hook(dct)  # Yes: then do it
        return dct  # No: just return the decoded dict

# return a long representations of the current time in milliseconds
def current_time_milliseconds():
    return int(round(time.time() * 1000))

# Generate a 'unique' md5 hash adding the current time in milliseconds
# to a string
def unique_md5(string: str):
    m = hashlib.md5()
    m.update(f'{current_time_milliseconds()}{string}'.encode())
    return m.hexdigest()

# Utility function to only execute and assigment to an object if the value from a reqparse.Namespace dict is not None
def _assign_if_something(obj: object, newdata: Namespace, key: str):
    value = newdata.get(key)
    if value is not None:
        obj.__setattr__(key, value)

# Apply filter restrictions
def restrict(query, filters, name, condition):
    f = filters.get(name)
    if f:
        query = query.filter(condition(f))
    return query

# Encrypt password
def sha1_pass(text: str):
    m = hashlib.sha1()
    m.update(text.encode('utf-8'))
    d = m.digest()
    t = ''
    for aux in d:
        c: int = aux & 0xff
        hs = '{:02x}'.format(c)
        t += hs
    return t

def check(permision):
    def wrfunc(fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            usuario = get_jwt_identity()
            if usuario is None:
                return {'message': 'No tiene permisos para realizar esta acción'}, 401
            claims = get_jwt_claims()
            if permision not in claims['permissions']:
                return {'message': 'No tiene permisos para realizar esta acción'}, 401
            return fn(*args, **kwargs)

        return wrapper

    return wrfunc

def paginated_results(query):
    pagination = request.args.get('pagination', 'true', str)
    jsondepth = request.args.get('jsondepth', 1, int)
    if pagination == 'true':
        paginated = query.paginate(page=request.args.get('page', 1, int))
        return {
            'page': paginated.page,
            'pages': paginated.pages,
            'items': [x.json(jsondepth) if jsondepth else x.json() for x in paginated.items]
        }
    else:
        return [x.json(jsondepth) if jsondepth else x.json() for x in query.all()]

# toHex = lambda x: "".join([hex(ord(c))[2:].zfill(2) for c in x])
# toHex = lambda x: "".join(["{:02X}".format(ord(c)) for c in x])
# Codificación hexadecimal
def to_hex(x):
    arr = []
    for c in x:
        try:
            coded = c.encode('latin_1')
            a = ord(c)
            b = hex(a)
            d = b[2:]
            e = d.zfill(2)
        except UnicodeEncodeError:
            byteString = c.encode('utf-8')
            e = ''.join('{:02x}'.format(x) for x in byteString)
        arr.append(e)
    return "".join(arr)

def check_request_sanity(values):
    if not 0 < values['port'] < 65536:
        return 'invalid port'
    if not values['left'] >= 0:
        return 'invalid amount remaining'
    if not values['downloaded'] >= 0:
        return 'invalid downloaded amount'
    if not values['uploaded'] >= 0:
        return 'invalid uploaded amount'
    if values['compact'] not in (0, 1):
        return 'compact must be 0 or 1'
    if not values['numwant'] >= 0:
        return 'invalid number of peers requested'
    if 'event' in values and values['event'] not in ('started', 'completed', 'stopped', 'update'):
        return 'invalid event'
    if values['peer_id'] is None:
        return 'peer id required'

    # all the data is good!
    return None

def parse_request(request):
    values = {}
    # request.encoding = 'iso-8859-1'

    parse_url = urlparse(request.url)
    query_dict = dict(parse_qsl(parse_url.query, encoding='latin_1'))

    # Get the interesting values from the request
    values['info_hash'] = to_hex(query_dict['info_hash'])
    values['peer_id'] = to_hex(query_dict['peer_id'])
    print("IPs -->", request.access_route, request.headers.getlist("X-Forwarded-For"))
    values['ip'] = request.environ['REMOTE_ADDR']  # ignore any sent ip address
    values['key'] = request.args.get('key', None)
    if request.args.get('event', None):  # Else Update
        values['event'] = request.args.get('event')

    # These are integer values and require special handling
    try:
        values['port'] = int(request.args.get('port', None))
        values['left'] = int(request.args.get('left', None))
        values['downloaded'] = int(request.args.get('downloaded', None))
        values['uploaded'] = int(request.args.get('uploaded', None))

        # Default values for parameters that arent required
        values['numwant'] = int(request.args.get('numwant', settings['numwant_default']))
        values['compact'] = int(request.args.get('compact', settings['compact_default']))
        values['no_peer_id'] = int(request.args.get('no_peer_id', settings['no_peer_id_default']))
    except ValueError:
        return None, 'error parsing integer field'
    except TypeError:
        return None, 'missing value for required field'

    return values, check_request_sanity(values)

def get_interval(peer):
    if peer.seeding:
        interval = settings['seed_interval']
    else:
        interval = settings['peer_interval']
    return interval

def tracker_error(value):
    return Response(bencode({'failure reason': "Error: %s" % value}), mimetype='text/plain')

def tracker_response(value):
    return Response(bencode(value), mimetype='text/plain')


import isolearn.keras as iso

import numpy as np

import pandas as pd

import scipy.sparse as sp
import scipy.io as spio

import matplotlib.pyplot as plt

import seaborn as sns

def plot_logo_w_hexamer_scores(pwm, iso_pred, cut_pred, hexamer_scores, cse_start_pos=70, annotate_peaks=None, sequence_template=None, figsize=(12, 3), width_ratios=[1, 7], logo_height=1.0, usage_unit='log', plot_start=0, plot_end=164, save_figs=False, fig_name=None, fig_dpi=300) :

    n_samples = pwm.shape[0]

    #Slice according to seq trim index
    pwm = pwm[:, plot_start: plot_end, :]
    cut_pred = cut_pred[:, plot_start: plot_end]
    sequence_template = sequence_template[plot_start: plot_end]

    iso_pred = np.mean(iso_pred, axis=0)
    cut_pred = np.mean(cut_pred, axis=0)
    pwm = np.sum(pwm, axis=0)

    pwm += 0.0001
    for j in range(0, pwm.shape[0]) :
        pwm[j, :] /= np.sum(pwm[j, :])

    entropy = np.zeros(pwm.shape)
    entropy[pwm > 0] = pwm[pwm > 0] * -np.log2(pwm[pwm > 0])
    entropy = np.sum(entropy, axis=1)
    conservation = 2 - entropy

    fig = plt.figure(figsize=figsize)

    gs = gridspec.GridSpec(2, 2, width_ratios=[width_ratios[0], width_ratios[-1]], height_ratios=[1, 1])

    ax0 = plt.subplot(gs[0, 0])
    ax1 = plt.subplot(gs[0, 1])
    ax2 = plt.subplot(gs[1, 0])
    ax3 = plt.subplot(gs[1, 1])

    plt.sca(ax0)
    plt.axis('off')
    plt.sca(ax2)
    plt.axis('off')

    annot_text = 'Samples = ' + str(int(n_samples))
    if usage_unit == 'log' :
        annot_text += '\nUsage = ' + str(round(np.log(iso_pred[0] / (1. - iso_pred[0])), 4))
    else :
        annot_text += '\nUsage = ' + str(round(iso_pred[0], 4))

    ax2.text(0.99, 0.5, annot_text, horizontalalignment='right', verticalalignment='center', transform=ax2.transAxes, color='black', fontsize=12, weight="bold")

    l2, = ax1.plot(np.arange(plot_end - plot_start), cut_pred, linewidth=3, linestyle='-', label='Predicted', color='red', alpha=0.7)

    if annotate_peaks is not None :
        objective_pos = 0
        if annotate_peaks == 'max' :
            objective_pos = np.argmax(cut_pred)
        else :
            objective_pos = annotate_peaks - plot_start

        text_x, text_y, ha = -30, -5, 'right'
        if objective_pos < 30 :
            text_x, text_y, ha = 30, -5, 'left'

        annot_text = '(CSE+' + str(objective_pos + plot_start - (cse_start_pos + 6) + 0) + ') ' + str(int(round(cut_pred[objective_pos] * 100, 0))) + '% Cleavage'
        ax1.annotate(annot_text, xy=(objective_pos, cut_pred[objective_pos]), xycoords='data', xytext=(text_x, text_y), ha=ha, fontsize=10, weight="bold", color='black', textcoords='offset points', arrowprops=dict(connectionstyle="arc3,rad=-.1", headlength=8, headwidth=8, shrink=0.15, width=1.5, color='black'))

    plt.sca(ax1)

    plt.xlim((0, plot_end - plot_start))
    #plt.ylim((0, 2))
    plt.xticks([], [])
    plt.yticks([], [])
    plt.legend(handles=[l2], fontsize=12, prop=dict(weight='bold'), frameon=False)
    plt.axis('off')

    height_base = (1.0 - logo_height) / 2.

    for j in range(0, pwm.shape[0]) :
        sort_index = np.argsort(pwm[j, :])

        for ii in range(0, 4) :
            i = sort_index[ii]

            nt_prob = pwm[j, i] * conservation[j]

            nt = ''
            if i == 0 :
                nt = 'A'
            elif i == 1 :
                nt = 'C'
            elif i == 2 :
                nt = 'G'
            elif i == 3 :
                nt = 'T'

            color = None
            if sequence_template[j] != 'N' :
                color = 'black'

            if ii == 0 :
                letterAt(nt, j + 0.5, height_base, nt_prob * logo_height, ax3, color=color)
            else :
                prev_prob = np.sum(pwm[j, sort_index[:ii]] * conservation[j]) * logo_height
                letterAt(nt, j + 0.5, height_base + prev_prob, nt_prob * logo_height, ax3, color=color)

            text_x, text_y, ha = -30, -5, 'right'
            if objective_pos < 30 :
                text_x, text_y, ha = 30, -5, 'left'

            if j < len(hexamer_scores) :
                #annot_text = str(hexamer_scores[j][0]) + " = " + str(round(hexamer_scores[j][1], 2))

                annot_text = str(hexamer_scores[j][1])
                ax3.text(j + 0.25, 2.0, annot_text, size=50, rotation=90., ha="left", va="bottom", fontsize=8, weight="bold", color='black')

    plt.sca(ax3)

    plt.xlim((0, plot_end - plot_start))
    plt.ylim((0, 2.5))
    plt.xticks([], [])
    plt.yticks([], [])
    plt.axis('off')
    ax3.axhline(y=0.01 + height_base, color='black', linestyle='-', linewidth=2)

    for axis in fig.axes :
        axis.get_xaxis().set_visible(False)
        axis.get_yaxis().set_visible(False)

    plt.tight_layout()

    if save_figs :
        plt.savefig(fig_name + '.png', transparent=True, dpi=fig_dpi)
        plt.savefig(fig_name + '.svg')
        plt.savefig(fig_name + '.eps')

    plt.show()
    #plt.close()

mer6_dict = {}
i = 0
for b1 in ['A', 'C', 'G', 'T'] :
    for b2 in ['A', 'C', 'G', 'T'] :
        for b3 in ['A', 'C', 'G', 'T'] :
            for b4 in ['A', 'C', 'G', 'T'] :
                for b5 in ['A', 'C', 'G', 'T'] :
                    for b6 in ['A', 'C', 'G', 'T'] :
                        mer6_dict[b1 + b2 + b3 + b4 + b5 + b6] = i
                        i += 1

def get_hexamer_pred_both_regions(seq) :
    preds = []
    for cell_line, w, w_0 in zip(['hek', 'hela', 'mcf7', 'cho'], [hek_both_regions_w, hela_both_regions_w, mcf7_both_regions_w, cho_both_regions_w], [hek_both_regions_w_0, hela_both_regions_w_0, mcf7_both_regions_w_0, cho_both_regions_w_0]) :
        pred = w_0

        region_1 = seq[5: 40]
        for j in range(0, len(region_1) - 5) :
            pred += w[mer6_dict[region_1[j: j+6]]]

        region_2 = seq[48: 83]
        for j in range(0, len(region_2) - 5) :
            pred += w[4096 + mer6_dict[region_2[j: j+6]]]

        preds.append(1. / (1. + np.exp(-pred)))

    return np.array(preds)

def get_hexamer_diff_scores_both_regions(seq, cell_1, cell_2) :
    scores = []

    w_dict = {
        'hek' : hek_both_regions_w,
        'hela' : hela_both_regions_w,
        'mcf7' : mcf7_both_regions_w,
        'cho' : cho_both_regions_w
    }

    w_cell_1 = w_dict[cell_1]
    w_cell_2 = w_dict[cell_2]

    scores.extend([('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', '')])

    region_1 = seq[5: 40]
    for j in range(0, len(region_1) - 5) :
        hexamer_score_cell_1 = w_cell_1[mer6_dict[region_1[j: j+6]]]
        hexamer_score_cell_2 = w_cell_2[mer6_dict[region_1[j: j+6]]]

        scores.append((region_1[j: j+6], str(round(hexamer_score_cell_1 - hexamer_score_cell_2, 2))))

    scores.extend([('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', '')])

    region_2 = seq[48: 83]
    for j in range(0, len(region_2) - 5) :
        hexamer_score_cell_1 = w_cell_1[4096 + mer6_dict[region_2[j: j+6]]]
        hexamer_score_cell_2 = w_cell_2[4096 + mer6_dict[region_2[j: j+6]]]

        scores.append((region_2[j: j+6], str(round(hexamer_score_cell_1 - hexamer_score_cell_2, 2))))

    return scores

def get_hexamer_preds_both_regions(seqs) :
    preds = np.zeros((len(seqs), 4))

    for i, seq in enumerate(seqs) :
        preds[i, :] = get_hexamer_pred_both_regions(seq)

    return preds

def get_hexamer_pred(seq) :
    preds = []
    for cell_line, w, w_0 in zip(['hek', 'hela', 'mcf7', 'cho'], [hek_w, hela_w, mcf7_w, cho_w], [hek_w_0, hela_w_0, mcf7_w_0, cho_w_0]) :
        pred = w_0

        region_1 = seq[5: 40]
        for j in range(0, len(region_1) - 5) :
            pred += w[mer6_dict[region_1[j: j+6]]]

        preds.append(1. / (1. + np.exp(-pred)))

    return np.array(preds)

def get_hexamer_diff_scores(seq, cell_1, cell_2) :
    scores = []

    w_dict = {
        'hek' : hek_w,
        'hela' : hela_w,
        'mcf7' : mcf7_w,
        'cho' : cho_w
    }

    w_cell_1 = w_dict[cell_1]
    w_cell_2 = w_dict[cell_2]

    scores.extend([('_', ''), ('_', ''), ('_', ''), ('_', ''), ('_', '')])
    region_1 = seq[5: 40]
    for j in range(0, len(region_1) - 5) :
        hexamer_score_cell_1 = w_cell_1[mer6_dict[region_1[j: j+6]]]
        hexamer_score_cell_2 = w_cell_2[mer6_dict[region_1[j: j+6]]]

        scores.append((region_1[j: j+6], str(round(hexamer_score_cell_1 - hexamer_score_cell_2, 2))))

    return scores

def get_hexamer_preds(seqs) :
    preds = np.zeros((len(seqs), 4))

    for i, seq in enumerate(seqs) :
        preds[i, :] = get_hexamer_pred(seq)

    return preds

def decode_onehot_consensus(onehot) :
    seq = ''

    for i in range(onehot.shape[0]) :
        max_j = np.argmax(onehot[i, :])

        if max_j == 0 :
            seq += 'A'
        elif max_j == 1 :
            seq += 'C'
        elif max_j == 2 :
            seq += 'G'
        elif max_j == 3 :
            seq += 'T'

    return seq

def decode_onehots_consensus(onehots) :
    seqs = [
        decode_onehot_consensus(onehots[i, :, :, 0]) for i in range(onehots.shape[0])
    ]

    return seqs