from functools import partial

import numpy as np

import mathlib
from facelib import FaceType
from interact import interact as io
from models import ModelBase
from nnlib import nnlib
from samplelib import *

#SAE - Styled AutoEncoder
class SAEHDModel(ModelBase):

    #override
    def onInitializeOptions(self, is_first_run, ask_override):
        yn_str = {True:'y',False:'n'}

        default_resolution = 128
        default_archi = 'df'
        default_face_type = 'f'

        if is_first_run:
            resolution = io.input_int("Resolution ( 64-256 ?:help skip:128) : ", default_resolution, help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.")
            resolution = np.clip (resolution, 64, 256)
            while np.modf(resolution / 16)[0] != 0.0:
                resolution -= 1
            self.options['resolution'] = resolution
            self.options['face_type'] = io.input_str ("Half, mid full, or full face? (h/mf/f, ?:help skip:f) : ", default_face_type, ['h','mf','f'], help_message="Half face has better resolution, but covers less area of cheeks. Mid face is 30% wider than half face.").lower()
        else:
            self.options['resolution'] = self.options.get('resolution', default_resolution)
            self.options['face_type'] = self.options.get('face_type', default_face_type)

        default_learn_mask = self.options.get('learn_mask', True)
        if is_first_run or ask_override:
            self.options['learn_mask'] = io.input_bool ( f"Learn mask? (y/n, ?:help skip:{yn_str[default_learn_mask]} ) : " , default_learn_mask, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted.")
        else:
            self.options['learn_mask'] = self.options.get('learn_mask', default_learn_mask)

        if (is_first_run or ask_override) and 'tensorflow' in self.device_config.backend:
            def_optimizer_mode = self.options.get('optimizer_mode', 1)
            self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.")
        else:
            self.options['optimizer_mode'] = self.options.get('optimizer_mode', 1)

        if is_first_run:
            self.options['archi'] = io.input_str ("AE architecture (df, liae ?:help skip:%s) : " % (default_archi) , default_archi, ['df','liae'], help_message="'df' keeps faces more natural. 'liae' can fix overly different face shapes.").lower() #-s version is slower, but has decreased change to collapse.
        else:
            self.options['archi'] = self.options.get('archi', default_archi)

        default_ae_dims = 256
        default_ed_ch_dims = 21

        if is_first_run:
            self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 )
            self.options['ed_ch_dims'] = np.clip ( io.input_int("Encoder/Decoder dims per channel (10-85 ?:help skip:%d) : " % (default_ed_ch_dims) , default_ed_ch_dims, help_message="More dims help to recognize more facial features and achieve sharper result, but require more VRAM. You can fine-tune model size to fit your GPU." ), 10, 85 )
        else:
            self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims)
            self.options['ed_ch_dims'] = self.options.get('ed_ch_dims', default_ed_ch_dims)

        default_true_face_training = self.options.get('true_face_training', False)
        default_face_style_power = self.options.get('face_style_power', 0.0)
        default_bg_style_power = self.options.get('bg_style_power', 0.0)

        if is_first_run or ask_override:
            default_random_warp = self.options.get('random_warp', True)
            self.options['random_warp'] = io.input_bool (f"Enable random warp of samples? ( y/n, ?:help skip:{yn_str[default_random_warp]}) : ", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness for less amount of iterations.")

            self.options['true_face_training'] = io.input_bool (f"Enable 'true face' training? (y/n, ?:help skip:{yn_str[default_true_face_training]}) : ", default_true_face_training, help_message="The result face will be more like src and will get extra sharpness. Enable it for last 10-20k iterations before conversion.")

            self.options['face_style_power'] = np.clip ( io.input_number("Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_face_style_power), default_face_style_power,
                                                                               help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes. Enabling this option increases the chance of model collapse."), 0.0, 100.0 )

            self.options['bg_style_power'] = np.clip ( io.input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_bg_style_power), default_bg_style_power,
                                                                               help_message="Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."), 0.0, 100.0 )

            default_ct_mode = self.options.get('ct_mode', 'none')
            self.options['ct_mode'] = io.input_str (f"Color transfer mode apply to src faceset. ( none/rct/lct/mkl/idt/sot, ?:help skip:{default_ct_mode}) : ", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.")

            if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301
                default_clipgrad = False if is_first_run else self.options.get('clipgrad', False)
                self.options['clipgrad'] = io.input_bool (f"Enable gradient clipping? (y/n, ?:help skip:{yn_str[default_clipgrad]}) : ", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.")
            else:
                self.options['clipgrad'] = False

        else:
            self.options['random_warp'] = self.options.get('random_warp', True)
            self.options['true_face_training'] = self.options.get('true_face_training', default_true_face_training)
            self.options['face_style_power'] = self.options.get('face_style_power', default_face_style_power)
            self.options['bg_style_power'] = self.options.get('bg_style_power', default_bg_style_power)
            self.options['ct_mode'] = self.options.get('ct_mode', 'none')
            self.options['clipgrad'] = self.options.get('clipgrad', False)

        if is_first_run:
            self.options['pretrain'] = io.input_bool ("Pretrain the model? (y/n, ?:help skip:n) : ", False, help_message="Pretrain the model with large amount of various faces. This technique may help to train the fake with overly different face shapes and light conditions of src/dst data. Face will be look more like a morphed. To reduce the morph effect, some model files will be initialized but not be updated after pretrain: LIAE: inter_AB.h5 DF: encoder.h5. The longer you pretrain the model the more morphed face will look. After that, save and run the training again.")
        else:
            self.options['pretrain'] = False

    #override
    def onInitialize(self):
        exec(nnlib.import_all(), locals(), globals())
        self.set_vram_batch_requirements({1.5:4,4:8})

        resolution = self.options['resolution']
        learn_mask = self.options['learn_mask']

        ae_dims = self.options['ae_dims']
        ed_ch_dims = self.options['ed_ch_dims']
        self.pretrain = self.options['pretrain'] = self.options.get('pretrain', False)
        if not self.pretrain:
            self.options.pop('pretrain')

        bgr_shape = (resolution, resolution, 3)
        mask_shape = (resolution, resolution, 1)

        self.true_face_training = self.options.get('true_face_training', False)
        masked_training = True

        class CommonModel(object):
            def downscale (self, dim, kernel_size=5, dilation_rate=1, use_activator=True):
                def func(x):
                    if not use_activator:
                        return SubpixelDownscaler()(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x))
                    else:
                        return SubpixelDownscaler()(LeakyReLU(0.1)(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x)))
                return func

            def upscale (self, dim, size=(2,2)):
                def func(x):
                    return SubpixelUpscaler(size=size)(LeakyReLU(0.1)(Conv2D(dim * np.prod(size) , kernel_size=3, strides=1, padding='same')(x)))
                return func

            def ResidualBlock(self, dim):
                def func(inp):
                    x = Conv2D(dim, kernel_size=3, padding='same')(inp)
                    x = LeakyReLU(0.2)(x)
                    x = Conv2D(dim, kernel_size=3, padding='same')(x)
                    x = Add()([x, inp])
                    x = LeakyReLU(0.2)(x)
                    return x
                return func

        class SAEDFModel(CommonModel):
            def __init__(self, resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask):
                super().__init__()
                self.learn_mask = learn_mask

                output_nc = 3
                bgr_shape = (resolution, resolution, output_nc)
                mask_shape = (resolution, resolution, 1)
                lowest_dense_res = resolution // 16
                e_dims = output_nc*e_ch_dims

                def enc_flow(e_ch_dims, ae_dims, lowest_dense_res):
                    dims = output_nc * e_ch_dims
                    if dims % 2 != 0:
                        dims += 1

                    def func(inp):
                        x = self.downscale(dims  , 3, 1 )(inp)
                        x = self.downscale(dims*2, 3, 1 )(x)
                        x = self.downscale(dims*4, 3, 1 )(x)
                        x0 = self.downscale(dims*8, 3, 1 )(x)

                        x = self.downscale(dims  , 5, 1 )(inp)
                        x = self.downscale(dims*2, 5, 1 )(x)
                        x = self.downscale(dims*4, 5, 1 )(x)
                        x1 = self.downscale(dims*8, 5, 1 )(x)

                        x = self.downscale(dims  , 5, 2 )(inp)
                        x = self.downscale(dims*2, 5, 2 )(x)
                        x = self.downscale(dims*4, 5, 2 )(x)
                        x2 = self.downscale(dims*8, 5, 2 )(x)

                        x = self.downscale(dims  , 7, 2 )(inp)
                        x = self.downscale(dims*2, 7, 2 )(x)
                        x = self.downscale(dims*4, 7, 2 )(x)
                        x3 = self.downscale(dims*8, 7, 2 )(x)

                        x = Concatenate()([x0,x1,x2,x3])

                        x = Dense(ae_dims)(Flatten()(x))
                        x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x)
                        x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x)
                        x = self.upscale(ae_dims)(x)
                        return x
                    return func

                def dec_flow(output_nc, d_ch_dims, is_mask=False):
                    dims = output_nc * d_ch_dims
                    if dims % 2 != 0:
                        dims += 1

                    def func(x):

                        for i in [8,4,2]:
                            x = self.upscale(dims*i)(x)

                            if not is_mask:
                                x0 = x
                                x = self.upscale( (dims*i)//2 )(x)
                                x = self.ResidualBlock( (dims*i)//2 )(x)
                                x = self.downscale( dims*i, use_activator=False ) (x)
                                x = Add()([x, x0])
                                x = LeakyReLU(0.2)(x)

                        return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x)

                    return func

                self.encoder = modelify(enc_flow(e_ch_dims, ae_dims, lowest_dense_res)) ( Input(bgr_shape) )

                sh = K.int_shape( self.encoder.outputs[0] )[1:]
                self.decoder_src = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) )
                self.decoder_dst = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) )

                if learn_mask:
                    self.decoder_srcm = modelify(dec_flow(1, d_ch_dims, is_mask=True)) ( Input(sh) )
                    self.decoder_dstm = modelify(dec_flow(1, d_ch_dims, is_mask=True)) ( Input(sh) )

                self.src_dst_trainable_weights = self.encoder.trainable_weights + self.decoder_src.trainable_weights + self.decoder_dst.trainable_weights

                if learn_mask:
                    self.src_dst_mask_trainable_weights = self.encoder.trainable_weights + self.decoder_srcm.trainable_weights + self.decoder_dstm.trainable_weights

                self.warped_src, self.warped_dst = Input(bgr_shape), Input(bgr_shape)
                self.target_src, self.target_dst = Input(bgr_shape), Input(bgr_shape)
                self.target_srcm, self.target_dstm = Input(mask_shape), Input(mask_shape)
                self.src_code, self.dst_code = self.encoder(self.warped_src), self.encoder(self.warped_dst)

                self.pred_src_src = self.decoder_src(self.src_code)
                self.pred_dst_dst = self.decoder_dst(self.dst_code)
                self.pred_src_dst = self.decoder_src(self.dst_code)

                if learn_mask:
                    self.pred_src_srcm = self.decoder_srcm(self.src_code)
                    self.pred_dst_dstm = self.decoder_dstm(self.dst_code)
                    self.pred_src_dstm = self.decoder_srcm(self.dst_code)

            def get_model_filename_list(self, exclude_for_pretrain=False):
                ar = []
                if not exclude_for_pretrain:
                    ar += [ [self.encoder, 'encoder.h5'] ]
                ar += [  [self.decoder_src, 'decoder_src.h5'],
                         [self.decoder_dst, 'decoder_dst.h5']  ]
                if self.learn_mask:
                    ar += [ [self.decoder_srcm, 'decoder_srcm.h5'],
                            [self.decoder_dstm, 'decoder_dstm.h5']  ]
                return ar

        class SAELIAEModel(CommonModel):
            def __init__(self, resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask):
                super().__init__()
                self.learn_mask = learn_mask

                output_nc = 3
                bgr_shape = (resolution, resolution, output_nc)
                mask_shape = (resolution, resolution, 1)

                lowest_dense_res = resolution // 16

                def enc_flow(e_ch_dims):
                    dims = output_nc*e_ch_dims
                    if dims % 2 != 0:
                        dims += 1

                    def func(inp):
                        x = self.downscale(dims  , 3, 1 )(inp)
                        x = self.downscale(dims*2, 3, 1 )(x)
                        x = self.downscale(dims*4, 3, 1 )(x)
                        x0 = self.downscale(dims*8, 3, 1 )(x)

                        x = self.downscale(dims  , 5, 1 )(inp)
                        x = self.downscale(dims*2, 5, 1 )(x)
                        x = self.downscale(dims*4, 5, 1 )(x)
                        x1 = self.downscale(dims*8, 5, 1 )(x)

                        x = self.downscale(dims  , 5, 2 )(inp)
                        x = self.downscale(dims*2, 5, 2 )(x)
                        x = self.downscale(dims*4, 5, 2 )(x)
                        x2 = self.downscale(dims*8, 5, 2 )(x)

                        x = self.downscale(dims  , 7, 2 )(inp)
                        x = self.downscale(dims*2, 7, 2 )(x)
                        x = self.downscale(dims*4, 7, 2 )(x)
                        x3 = self.downscale(dims*8, 7, 2 )(x)

                        x = Concatenate()([x0,x1,x2,x3])

                        x = Flatten()(x)
                        return x
                    return func

                def inter_flow(lowest_dense_res, ae_dims):
                    def func(x):
                        x = Dense(ae_dims)(x)
                        x = Dense(lowest_dense_res * lowest_dense_res * ae_dims*2)(x)
                        x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims*2))(x)
                        x = self.upscale(ae_dims*2)(x)
                        return x
                    return func

                def dec_flow(output_nc, d_ch_dims, is_mask=False):
                    dims = output_nc * d_ch_dims
                    if dims % 2 != 0:
                        dims += 1

                    def func(x):

                        for i in [8,4,2]:
                            x = self.upscale(dims*i)(x)

                            if not is_mask:
                                x0 = x
                                x = self.upscale( (dims*i)//2 )(x)
                                x = self.ResidualBlock( (dims*i)//2 )(x)
                                x = self.downscale( dims*i, use_activator=False ) (x)
                                x = Add()([x, x0])
                                x = LeakyReLU(0.2)(x)

                        return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x)

                    return func

                self.encoder = modelify(enc_flow(e_ch_dims)) ( Input(bgr_shape) )

                sh = K.int_shape( self.encoder.outputs[0] )[1:]
                self.inter_B = modelify(inter_flow(lowest_dense_res, ae_dims)) ( Input(sh) )
                self.inter_AB = modelify(inter_flow(lowest_dense_res, ae_dims)) ( Input(sh) )

                sh = np.array(K.int_shape( self.inter_B.outputs[0] )[1:])*(1,1,2)
                self.decoder = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) )

                if learn_mask:
                    self.decoderm = modelify(dec_flow(1, d_ch_dims, is_mask=True)) ( Input(sh) )

                self.src_dst_trainable_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoder.trainable_weights

                if learn_mask:
                    self.src_dst_mask_trainable_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoderm.trainable_weights

                self.warped_src, self.warped_dst = Input(bgr_shape), Input(bgr_shape)
                self.target_src, self.target_dst = Input(bgr_shape), Input(bgr_shape)
                self.target_srcm, self.target_dstm = Input(mask_shape), Input(mask_shape)

                warped_src_code = self.encoder (self.warped_src)
                warped_src_inter_AB_code = self.inter_AB (warped_src_code)
                self.src_code = Concatenate()([warped_src_inter_AB_code,warped_src_inter_AB_code])

                warped_dst_code = self.encoder (self.warped_dst)
                warped_dst_inter_B_code = self.inter_B (warped_dst_code)
                warped_dst_inter_AB_code = self.inter_AB (warped_dst_code)
                self.dst_code = Concatenate()([warped_dst_inter_B_code,warped_dst_inter_AB_code])

                src_dst_code = Concatenate()([warped_dst_inter_AB_code,warped_dst_inter_AB_code])

                self.pred_src_src = self.decoder(self.src_code)
                self.pred_dst_dst = self.decoder(self.dst_code)
                self.pred_src_dst = self.decoder(src_dst_code)

                if learn_mask:
                    self.pred_src_srcm = self.decoderm(self.src_code)
                    self.pred_dst_dstm = self.decoderm(self.dst_code)
                    self.pred_src_dstm = self.decoderm(src_dst_code)

            def get_model_filename_list(self, exclude_for_pretrain=False):
                ar = [ [self.encoder, 'encoder.h5'],
                       [self.inter_B, 'inter_B.h5'] ]

                if not exclude_for_pretrain:
                    ar += [ [self.inter_AB, 'inter_AB.h5'] ]

                ar += [  [self.decoder, 'decoder.h5']  ]

                if self.learn_mask:
                    ar += [ [self.decoderm, 'decoderm.h5'] ]

                return ar

        if 'df' in self.options['archi']:
            self.model = SAEDFModel (resolution, ae_dims, ed_ch_dims, ed_ch_dims, learn_mask)
        elif 'liae' in self.options['archi']:
            self.model = SAELIAEModel (resolution, ae_dims, ed_ch_dims, ed_ch_dims, learn_mask)

        self.opt_dis_model = []

        if self.true_face_training:
            def dis_flow(ndf=256):
                def func(x):
                    x, = x

                    code_res = K.int_shape(x)[1]

                    x = Conv2D( ndf, 4, strides=2, padding='valid')( ZeroPadding2D(1)(x) )
                    x = LeakyReLU(0.1)(x)

                    x = Conv2D( ndf*2, 3, strides=2, padding='valid')( ZeroPadding2D(1)(x) )
                    x = LeakyReLU(0.1)(x)

                    if code_res > 8:
                        x = Conv2D( ndf*4, 3, strides=2, padding='valid')( ZeroPadding2D(1)(x) )
                        x = LeakyReLU(0.1)(x)

                    if code_res > 16:
                        x = Conv2D( ndf*8, 3, strides=2, padding='valid')( ZeroPadding2D(1)(x) )
                        x = LeakyReLU(0.1)(x)

                    if code_res > 32:
                        x = Conv2D( ndf*8, 3, strides=2, padding='valid')( ZeroPadding2D(1)(x) )
                        x = LeakyReLU(0.1)(x)

                    return Conv2D( 1, 1, strides=1, padding='valid', activation='sigmoid')(x)
                return func

            sh = [ Input( K.int_shape(self.model.src_code)[1:] ) ]
            self.dis = modelify(dis_flow()) (sh)

            self.opt_dis_model = [ (self.dis, 'dis.h5') ]

        loaded, not_loaded = [], self.model.get_model_filename_list()+self.opt_dis_model
        if not self.is_first_run():
            loaded, not_loaded = self.load_weights_safe(not_loaded)

        CA_models = [ model for model, _ in not_loaded ]

        self.CA_conv_weights_list = []
        for model in CA_models:
            for layer in model.layers:
                if type(layer) == keras.layers.Conv2D:
                    self.CA_conv_weights_list += [layer.weights[0]] #- is Conv2D kernel_weights

        target_srcm = gaussian_blur( max(1, resolution // 32) )(self.model.target_srcm)
        target_dstm = gaussian_blur( max(1, resolution // 32) )(self.model.target_dstm)

        target_src_masked = self.model.target_src*target_srcm
        target_dst_masked = self.model.target_dst*target_dstm
        target_dst_anti_masked = self.model.target_dst*(1.0 - target_dstm)

        target_src_masked_opt = target_src_masked if masked_training else self.model.target_src
        target_dst_masked_opt = target_dst_masked if masked_training else self.model.target_dst

        pred_src_src_masked_opt = self.model.pred_src_src*target_srcm if masked_training else self.model.pred_src_src
        pred_dst_dst_masked_opt = self.model.pred_dst_dst*target_dstm if masked_training else self.model.pred_dst_dst

        psd_target_dst_masked = self.model.pred_src_dst*target_dstm
        psd_target_dst_anti_masked = self.model.pred_src_dst*(1.0 - target_dstm)

        if self.is_training_mode:
            self.src_dst_opt      = RMSprop(lr=5e-5, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1)
            self.src_dst_mask_opt = RMSprop(lr=5e-5, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1)
            self.D_opt            = RMSprop(lr=5e-5, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1)

            src_loss =  K.mean ( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_opt, pred_src_src_masked_opt) )
            src_loss += K.mean ( 10*K.square( target_src_masked_opt - pred_src_src_masked_opt ) )

            face_style_power = self.options['face_style_power'] / 100.0
            if face_style_power != 0:
                src_loss += style_loss(gaussian_blur_radius=resolution//16, loss_weight=face_style_power, wnd_size=0)( psd_target_dst_masked, target_dst_masked )

            bg_style_power = self.options['bg_style_power'] / 100.0
            if bg_style_power != 0:
                src_loss += K.mean( (10*bg_style_power)*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( psd_target_dst_anti_masked, target_dst_anti_masked ))
                src_loss += K.mean( (10*bg_style_power)*K.square( psd_target_dst_anti_masked - target_dst_anti_masked ))

            dst_loss =  K.mean( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)(target_dst_masked_opt, pred_dst_dst_masked_opt) )
            dst_loss += K.mean( 10*K.square( target_dst_masked_opt - pred_dst_dst_masked_opt ) )

            G_loss = src_loss+dst_loss

            if self.true_face_training:
                def DLoss(labels,logits):
                    return K.mean(K.binary_crossentropy(labels,logits))

                src_code_d = self.dis( self.model.src_code )
                src_code_d_ones = K.ones_like(src_code_d)
                src_code_d_zeros = K.zeros_like(src_code_d)
                dst_code_d = self.dis( self.model.dst_code )
                dst_code_d_ones = K.ones_like(dst_code_d)
                G_loss += 0.01*DLoss(src_code_d_ones, src_code_d)

                loss_D = (DLoss(dst_code_d_ones , dst_code_d) + \
                          DLoss(src_code_d_zeros, src_code_d) ) * 0.5

                self.D_train = K.function ([self.model.warped_src, self.model.warped_dst],[loss_D], self.D_opt.get_updates(loss_D, self.dis.trainable_weights) )

            self.src_dst_train = K.function ([self.model.warped_src, self.model.warped_dst, self.model.target_src, self.model.target_srcm, self.model.target_dst, self.model.target_dstm],
                                             [src_loss,dst_loss],
                                             self.src_dst_opt.get_updates( G_loss, self.model.src_dst_trainable_weights)
                                             )

            if self.options['learn_mask']:
                src_mask_loss = K.mean(K.square(self.model.target_srcm-self.model.pred_src_srcm))
                dst_mask_loss = K.mean(K.square(self.model.target_dstm-self.model.pred_dst_dstm))
                self.src_dst_mask_train = K.function ([self.model.warped_src, self.model.warped_dst, self.model.target_srcm, self.model.target_dstm],[src_mask_loss, dst_mask_loss], self.src_dst_mask_opt.get_updates(src_mask_loss+dst_mask_loss, self.model.src_dst_mask_trainable_weights ) )

            if self.options['learn_mask']:
                self.AE_view = K.function ([self.model.warped_src, self.model.warped_dst], [self.model.pred_src_src, self.model.pred_dst_dst, self.model.pred_dst_dstm, self.model.pred_src_dst, self.model.pred_src_dstm])
            else:
                self.AE_view = K.function ([self.model.warped_src, self.model.warped_dst], [self.model.pred_src_src, self.model.pred_dst_dst, self.model.pred_src_dst ])

        else:
            if self.options['learn_mask']:
                self.AE_convert = K.function ([self.model.warped_dst],[ self.model.pred_src_dst, self.model.pred_dst_dstm, self.model.pred_src_dstm ])
            else:
                self.AE_convert = K.function ([self.model.warped_dst],[ self.model.pred_src_dst ])

        if self.is_training_mode:
            t = SampleProcessor.Types

            if self.options['face_type'] == 'h':
                face_type = t.FACE_TYPE_HALF
            elif self.options['face_type'] == 'mf':
                face_type = t.FACE_TYPE_MID_FULL
            elif self.options['face_type'] == 'f':
                face_type = t.FACE_TYPE_FULL

            t_mode_bgr = t.MODE_BGR if not self.pretrain else t.MODE_BGR_SHUFFLE

            training_data_src_path = self.training_data_src_path
            training_data_dst_path = self.training_data_dst_path
            sort_by_yaw = self.sort_by_yaw

            if self.pretrain and self.pretraining_data_path is not None:
                training_data_src_path = self.pretraining_data_path
                training_data_dst_path = self.pretraining_data_path
                sort_by_yaw = False

            t_img_warped = t.IMG_WARPED_TRANSFORMED if self.options['random_warp'] else t.IMG_TRANSFORMED

            self.set_training_data_generators ([
                    SampleGeneratorFace(training_data_src_path, sort_by_yaw_target_samples_path=training_data_dst_path if sort_by_yaw else None,
                                                                random_ct_samples_path=training_data_dst_path if self.options['ct_mode'] != 'none' else None,
                                                                debug=self.is_debug(), batch_size=self.batch_size,
                        sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
                        output_sample_types = [ {'types' : (t_img_warped, face_type, t_mode_bgr), 'resolution':resolution, 'ct_mode': self.options['ct_mode'] },
                                                {'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution, 'ct_mode': self.options['ct_mode'] },
                                                {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution } ]
                                              ),

                    SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
                        sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ),
                        output_sample_types = [ {'types' : (t_img_warped, face_type, t_mode_bgr), 'resolution':resolution},
                                                {'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution},
                                                {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution} ])
                             ])

    #override
    def get_model_filename_list(self):
        return self.model.get_model_filename_list ( exclude_for_pretrain=(self.pretrain and self.iter != 0) ) +self.opt_dis_model

    #override
    def onSave(self):
        self.save_weights_safe( self.get_model_filename_list()+self.opt_dis_model )

    #override
    def on_success_train_one_iter(self):
        if len(self.CA_conv_weights_list) != 0:
            exec(nnlib.import_all(), locals(), globals())
            CAInitializerMP ( self.CA_conv_weights_list )
            self.CA_conv_weights_list = []

    #override
    def onTrainOneIter(self, generators_samples, generators_list):
        warped_src, target_src, target_srcm = generators_samples[0]
        warped_dst, target_dst, target_dstm = generators_samples[1]

        feed = [warped_src, warped_dst, target_src, target_srcm, target_dst, target_dstm]

        src_loss, dst_loss, = self.src_dst_train (feed)

        if self.true_face_training:
            self.D_train([warped_src, warped_dst])

        if self.options['learn_mask']:
            feed = [ warped_src, warped_dst, target_srcm, target_dstm ]
            src_mask_loss, dst_mask_loss, = self.src_dst_mask_train (feed)

        return ( ('src_loss', src_loss), ('dst_loss', dst_loss), )

    #override
    def onGetPreview(self, sample):
        test_S   = sample[0][1][0:4] #first 4 samples
        test_S_m = sample[0][2][0:4] #first 4 samples
        test_D   = sample[1][1][0:4]
        test_D_m = sample[1][2][0:4]

        if self.options['learn_mask']:
            S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ]
            DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ]
        else:
            S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ]

        result = []
        st = []
        for i in range(len(test_S)):
            ar = S[i], SS[i], D[i], DD[i], SD[i]

            st.append ( np.concatenate ( ar, axis=1) )

        result += [ ('SAEHD', np.concatenate (st, axis=0 )), ]

        if self.options['learn_mask']:
            st_m = []
            for i in range(len(test_S)):
                ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i])
                st_m.append ( np.concatenate ( ar, axis=1) )

            result += [ ('SAEHD masked', np.concatenate (st_m, axis=0 )), ]

        return result

    def predictor_func (self, face=None, dummy_predict=False):
        if dummy_predict:
            self.AE_convert ([ np.zeros ( (1, self.options['resolution'], self.options['resolution'], 3), dtype=np.float32 ) ])
        else:
            if self.options['learn_mask']:
                bgr, mask_dst_dstm, mask_src_dstm = self.AE_convert ([face[np.newaxis,...]])
                mask = mask_dst_dstm[0] * mask_src_dstm[0]
                return bgr[0], mask[...,0]
            else:
                bgr, = self.AE_convert ([face[np.newaxis,...]])
                return bgr[0]

    #override
    def get_ConverterConfig(self):
        if self.options['face_type'] == 'h':
            face_type = FaceType.HALF
        elif self.options['face_type'] == 'mf':
            face_type = FaceType.MID_FULL
        elif self.options['face_type'] == 'f':
            face_type = FaceType.FULL

        import converters
        return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=face_type,
                                     default_mode = 'overlay' if self.options['ct_mode'] != 'none' or self.options['face_style_power'] or self.options['bg_style_power'] else 'seamless',
                                     clip_hborder_mask_per=0.0625 if (face_type != FaceType.HALF) else 0,
                                    )

Model = SAEHDModel


# coding: utf-8
from __future__ import unicode_literals

import re

from .common import InfoExtractor
from ..utils import (
    determine_ext,
    int_or_none,
    float_or_none,
    js_to_json,
    parse_iso8601,
    remove_end,
)

class TV2IE(InfoExtractor):
    _VALID_URL = r'https?://(?:www\.)?tv2\.no/v/(?P<id>\d+)'
    _TEST = {
        'url': 'http://www.tv2.no/v/916509/',
        'info_dict': {
            'id': '916509',
            'ext': 'mp4',
            'title': 'Se Frode Gryttens hyllest av Steven Gerrard',
            'description': 'TV 2 Sportens huspoet tar avskjed med Liverpools kaptein Steven Gerrard.',
            'timestamp': 1431715610,
            'upload_date': '20150515',
            'duration': 156.967,
            'view_count': int,
            'categories': list,
        },
        'params': {
            # m3u8 download
            'skip_download': True,
        },
    }

    def _real_extract(self, url):
        video_id = self._match_id(url)

        formats = []
        format_urls = []
        for protocol in ('HDS', 'HLS'):
            data = self._download_json(
                'http://sumo.tv2.no/api/web/asset/%s/play.json?protocol=%s&videoFormat=SMIL+ISMUSP' % (video_id, protocol),
                video_id, 'Downloading play JSON')['playback']
            for item in data['items']['item']:
                video_url = item.get('url')
                if not video_url or video_url in format_urls:
                    continue
                format_id = '%s-%s' % (protocol.lower(), item.get('mediaFormat'))
                if not self._is_valid_url(video_url, video_id, format_id):
                    continue
                format_urls.append(video_url)
                ext = determine_ext(video_url)
                if ext == 'f4m':
                    formats.extend(self._extract_f4m_formats(
                        video_url, video_id, f4m_id=format_id, fatal=False))
                elif ext == 'm3u8':
                    formats.extend(self._extract_m3u8_formats(
                        video_url, video_id, 'mp4', entry_protocol='m3u8_native',
                        m3u8_id=format_id, fatal=False))
                elif ext == 'ism' or video_url.endswith('.ism/Manifest'):
                    pass
                else:
                    formats.append({
                        'url': video_url,
                        'format_id': format_id,
                        'tbr': int_or_none(item.get('bitrate')),
                        'filesize': int_or_none(item.get('fileSize')),
                    })
        self._sort_formats(formats)

        asset = self._download_json(
            'http://sumo.tv2.no/api/web/asset/%s.json' % video_id,
            video_id, 'Downloading metadata JSON')['asset']

        title = asset['title']
        description = asset.get('description')
        timestamp = parse_iso8601(asset.get('createTime'))
        duration = float_or_none(asset.get('accurateDuration') or asset.get('duration'))
        view_count = int_or_none(asset.get('views'))
        categories = asset.get('keywords', '').split(',')

        thumbnails = [{
            'id': thumbnail.get('@type'),
            'url': thumbnail.get('url'),
        } for _, thumbnail in asset.get('imageVersions', {}).items()]

        return {
            'id': video_id,
            'url': video_url,
            'title': title,
            'description': description,
            'thumbnails': thumbnails,
            'timestamp': timestamp,
            'duration': duration,
            'view_count': view_count,
            'categories': categories,
            'formats': formats,
        }

class TV2ArticleIE(InfoExtractor):
    _VALID_URL = r'https?://(?:www\.)?tv2\.no/(?:a|\d{4}/\d{2}/\d{2}(/[^/]+)+)/(?P<id>\d+)'
    _TESTS = [{
        'url': 'http://www.tv2.no/2015/05/16/nyheter/alesund/krim/pingvin/6930542',
        'info_dict': {
            'id': '6930542',
            'title': 'Russen hetses etter pingvintyveri - innrømmer å ha åpnet luken på buret',
            'description': 'md5:339573779d3eea3542ffe12006190954',
        },
        'playlist_count': 2,
    }, {
        'url': 'http://www.tv2.no/a/6930542',
        'only_matching': True,
    }]

    def _real_extract(self, url):
        playlist_id = self._match_id(url)

        webpage = self._download_webpage(url, playlist_id)

        # Old embed pattern (looks unused nowadays)
        assets = re.findall(r'data-assetid=["\'](\d+)', webpage)

        if not assets:
            # New embed pattern
            for v in re.findall(r'TV2ContentboxVideo\(({.+?})\)', webpage):
                video = self._parse_json(
                    v, playlist_id, transform_source=js_to_json, fatal=False)
                if not video:
                    continue
                asset = video.get('assetId')
                if asset:
                    assets.append(asset)

        entries = [
            self.url_result('http://www.tv2.no/v/%s' % asset_id, 'TV2')
            for asset_id in assets]

        title = remove_end(self._og_search_title(webpage), ' - TV2.no')
        description = remove_end(self._og_search_description(webpage), ' - TV2.no')

        return self.playlist_result(entries, playlist_id, title, description)


# vim: set ts=4 sts=4 sw=4 et fileencoding=utf-8:
#
# database related functions for ardj.
#
# ardj is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# ardj is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""ARDJ, an artificial DJ.

This module contains the database related code.
"""

import logging
import os
import random
import re
import sys
import time
import traceback
import urllib.request
import urllib.parse
import urllib.error

try:
    from sqlite3 import dbapi2 as sqlite
    from sqlite3 import OperationalError
except ImportError:
    print('Please install pysqlite2.', file=sys.stderr)
    sys.exit(13)

import ardj.settings
import ardj.scrobbler
import ardj.tags
import ardj.util

RECENT_SECONDS = 2 * 3600

SQL_INIT = [
    # playlists
    "CREATE TABLE IF NOT EXISTS playlists (id INTEGER PRIMARY KEY, name TEXT, last_played INTEGER);",

    # tracks
    "CREATE TABLE IF NOT EXISTS tracks (id INTEGER PRIMARY KEY, owner TEXT, filename TEXT, artist TEXT, title TEXT, length INTEGER, weight REAL, real_weight REAL, count INTEGER, last_played INTEGER, image TEXT, download TEXT);",
    "CREATE INDEX IF NOT EXISTS idx_tracks_owner ON tracks (owner);",
    "CREATE INDEX IF NOT EXISTS idx_tracks_last ON tracks (last_played);",
    "CREATE INDEX IF NOT EXISTS idx_tracks_count ON tracks (count);",
    "CREATE INDEX IF NOT EXISTS idx_tracks_weight ON tracks (weight);",
    "CREATE INDEX IF NOT EXISTS idx_tracks_real_weight ON tracks (real_weight);",
    "CREATE TABLE IF NOT EXISTS queue (id INTEGER PRIMARY KEY, track_id INTEGER, owner TEXT);",

    # urgent (manual) playlist
    "CREATE TABLE IF NOT EXISTS urgent_playlists (labels TEXT, expires INTEGER);",
    "CREATE INDEX IF NOT EXISTS urgent_playlists_expires ON urgent_playlists (expires);",

    # labels
    "CREATE TABLE IF NOT EXISTS labels (track_id INTEGER NOT NULL, email TEXT NOT NULL, label TEXT NOT NULL);",
    "CREATE INDEX IF NOT EXISTS idx_labels_track_id ON labels (track_id);",
    "CREATE INDEX IF NOT EXISTS idx_labels_email ON labels (email);",
    "CREATE INDEX IF NOT EXISTS idx_labels_label ON labels (label);",

    # voting
    "CREATE TABLE IF NOT EXISTS votes (track_id INTEGER NOT NULL, email TEXT NOT NULL, vote INTEGER, weight REAL, ts INTEGER);",
    "CREATE INDEX IF NOT EXISTS idx_votes_track_id ON votes (track_id);",
    "CREATE INDEX IF NOT EXISTS idx_votes_email ON votes (email);",
    "CREATE INDEX IF NOT EXISTS idx_votes_ts ON votes (ts);",

    # karma
    "CREATE TABLE IF NOT EXISTS karma (email TEXT, weight REAL);",
    "CREATE INDEX IF NOT EXISTS idx_karma_email ON karma (email);",

    # play log
    "CREATE TABLE IF NOT EXISTS playlog (ts INTEGER NOT NULL, track_id INTEGER NOT NULL, listeners INTEGER NOT NULL, lastfm INTEGER NOT NULL DEFAULT 0, librefm INTEGER NOT NULL DEFAULT 0);",
    "CREATE INDEX IF NOT EXISTS idx_playlog_ts ON playlog (ts);",
    "CREATE INDEX IF NOT EXISTS idx_playlog_track_id ON playlog (track_id);",

    # chat messages
    "CREATE TABLE IF NOT EXISTS jabber_messages (id INTEGER PRIMARY KEY, re TEXT, message TEXT);",

    # music downloads
    "CREATE TABLE IF NOT EXISTS download_queue (artist TEXT PRIMARY KEY, owner TEXT);",

    # web authentication
    "CREATE TABLE IF NOT EXISTS tokens (token TEXT PRIMARY KEY NOT NULL, login TEXT NOT NULL, login_type TEXT NOT NULL, active INTEGER NOT NULL DEFAULT 0);",
]

class Model(dict):
    table_name = None
    fields = ()
    key_name = None

    @classmethod
    def get_by_id(cls, id):
        fields_sql = ", ".join(cls.fields)
        sql = "SELECT %s FROM %s WHERE %s = ?" % (
            fields_sql, cls.table_name, cls.key_name)
        row = fetchone(sql, (id, ))
        if row is not None:
            return cls._from_row(row)

    @classmethod
    def find_all(cls):
        fields_sql = ", ".join(cls.fields)
        sql = "SELECT %s FROM %s" % (fields_sql, cls.table_name)
        return cls._fetch_rows(sql, ())

    @classmethod
    def _fetch_rows(cls, sql, params):
        rows = fetch(sql, params)
        return [cls._from_row(row) for row in rows]

    @classmethod
    def _from_row(cls, row):
        pairs = [(name, row[idx]) for idx, name in enumerate(cls.fields)]
        return cls(dict(pairs))

    @classmethod
    def _fields_sql(cls):
        return ", ".join(cls.fields)

    @classmethod
    def delete_all(cls):
        """Deletes all records from the table."""
        return execute("DELETE FROM %s" % cls.table_name, ())

    def delete(self):
        sql = "DELETE FROM %s WHERE %s = ?" % (self.table_name, self.key_name)
        execute(sql, (self[self.key_name], ))
        self[self.key_name] = None

    def put(self, force_insert=False):
        if self.get(self.key_name) is None or force_insert:
            return self._insert(force_insert)
        return self._update()

    def _insert(self, with_key=False):
        if with_key:
            fields = self.fields
        else:
            fields = [f for f in self.fields if f != self.key_name]
        fields_sql = ", ".join(fields)
        params_sql = ", ".join(["?"] * len(fields))

        sql = "INSERT INTO %s (%s) VALUES (%s)" % (
            self.table_name, fields_sql, params_sql)
        params = [self.get(field) for field in fields]

        self[self.key_name] = execute(sql, params)
        return self[self.key_name]

    def _update(self):
        fields = [f for f in self.fields if f != self.key_name]
        fields_sql = ", ".join(["%s = ?" % field for field in fields])

        sql = "UPDATE %s SET %s WHERE %s = ?" % (
            self.table_name, fields_sql, self.key_name)
        params = [self.get(field) for field in fields] + [self[self.key_name]]

        return execute(sql, params)

class Message(Model):
    """Represents an outgoing XMPP message."""
    table_name = "jabber_messages"
    fields = "id", "message", "re"
    key_name = "id"

class DownloadRequest(Model):
    table_name = "download_queue"
    fields = "artist", "owner"

    @classmethod
    def find_by_artist(cls, artist):
        sql = "SELECT %s FROM %s WHERE artist = ?" % (
            cls._fields_sql(), cls.table_name)
        return cls._fetch_rows(sql, (artist, ))

    @classmethod
    def get_one(cls):
        """Returns one random download ticket."""
        sql = "SELECT %s FROM %s LIMIT 1" % (cls._fields_sql(), cls.table_name)
        rows = cls._fetch_rows(sql, ())
        if rows:
            return rows[0]

class Vote(Model):
    """Stores information about votes."""
    table_name = "votes"
    fields = ("track_id", "email", "vote", "ts")

    def get_date(self):
        """Returns a formatted date for this vote."""
        return time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime(self["ts"]))

class Track(Model):
    """Stores information about a track."""
    table_name = "tracks"
    fields = (
        "id",
        "artist",
        "title",
        "filename",
        "length",
        "weight",
        "real_weight",
        "count",
        "last_played",
        "owner",
        "image",
        "download")
    key_name = "id"

    @classmethod
    def find_all(cls, deleted=False):
        """Returns all tracks with positive weight."""
        sql = "SELECT %s FROM %s" % (cls._fields_sql(), cls.table_name)
        if not deleted:
            sql += " WHERE weight > 0"
        return cls._fetch_rows(sql, ())

    @classmethod
    def find_by_tag(cls, tag):
        """Returns tracks that have a specific tag."""
        sql = "SELECT %s FROM %s WHERE weight > 0 AND id IN (SELECT track_id FROM labels WHERE label = ?)" % (
            cls._fields_sql(), cls.table_name)
        return cls._fetch_rows(sql, (tag, ))

    @classmethod
    def find_by_url(cls, url):
        """Returns all tracks with the specified download URL."""
        sql = "SELECT %s FROM %s WHERE `download` = ?" % (
            cls._fields_sql(), cls.table_name)
        return cls._fetch_rows(sql, (url, ))

    @classmethod
    def find_active(cls):
        """Returns tracks which weren't deleted."""
        sql = "SELECT %s FROM %s WHERE `weight` > 0" % (
            cls._fields_sql(), cls.table_name)
        return cls._fetch_rows(sql, ())

    @classmethod
    def find_recently_played(cls, count=50):
        """Returns a list of recently played tracks."""
        sql = "SELECT %s FROM %s ORDER BY `last_played` DESC LIMIT %u" % (
            cls._fields_sql(), cls.table_name, count)
        return cls._fetch_rows(sql, ())

    @classmethod
    def get_artist_names(cls):
        """Returns all artist names."""
        return fetchcol("SELECT DISTINCT artist FROM %s" % self.table_name)

    @classmethod
    def rename_artist(cls, old_name, new_name):
        """Renames an artist."""
        sql = "UPDATE %s SET artist = ? WHERE artist = ?" % (
            self.table_name, new_name, old_name)
        execute(sql, ())

    @classmethod
    def find_without_lastfm_tags(cls):
        sql = "SELECT %s FROM %s WHERE weight > 0 AND id NOT IN (SELECT track_id FROM labels WHERE label LIKE 'lastfm:%%') ORDER BY id" % (
            cls._fields_sql(), cls.table_name)
        return cls._fetch_rows(sql, ())

    @classmethod
    def find_tags(cls, min_count=5, cents=100):
        sql = "SELECT label, COUNT(*) FROM labels WHERE track_id IN (SELECT id FROM tracks WHERE weight > 0) GROUP BY label ORDER BY label"
        rows = [
            row for row in fetch(sql) if row[1] >= min_count and ":" not in row[0] and len(
                row[0]) > 1]
        return rows

    def get_labels(self):
        if not self.get(self.key_name):
            return []
        if "labels" not in self:
            self["labels"] = fetchcol(
                "SELECT label FROM labels WHERE track_id = ?", (self[self.key_name], ))
        return self["labels"] or []

    def set_labels(self, labels):
        if type(labels) != list:
            raise TypeError("Labels must be a list.")

        execute("DELETE FROM labels WHERE track_id = ?", (self["id"], ))
        for tag in list(set(labels)):
            execute(
                "INSERT INTO labels (track_id, label, email) VALUES (?, ?, ?)",
                (self["id"],
                 tag,
                 "unknown",
                 ))

        logging.debug(
            ("New labels for track %u: %s" %
             (self["id"], labels)).encode("utf-8"))

    def set_image(self, url):
        execute("UPDATE tracks SET image = ? WHERE id = ?",
                (url, self["id"], ))

    def set_download(self, url):
        execute("UPDATE tracks SET download = ? WHERE id = ?",
                (url, self["id"], ))

    def get_artist_url(self):
        if "lastfm:noartist" in self.get_labels():
            return None

        def q(v): return urllib.parse.quote(v.encode("utf-8"))
        return "http://www.last.fm/music/%s" % q(self["artist"])

    def get_track_url(self):
        if "lastfm:notfound" in self.get_labels():
            return None

        def q(v): return urllib.parse.quote(v.encode("utf-8"))
        return "http://www.last.fm/music/%s/_/%s" % (
            q(self["artist"]), q(self["title"]))

    @classmethod
    def get_average_length(cls):
        """Returns average track length in minutes."""
        s_prc = s_qty = 0.0
        for prc, qty in fetch(
                "SELECT ROUND(length / 60) AS r, COUNT(*) FROM tracks GROUP BY r"):
            s_prc += prc * qty
            s_qty += qty
        return int(s_prc / s_qty * 60 * 1.5)

    def get_votes(self):
        """Returns track votes."""
        result = {}
        for email, vote in fetch(
                "SELECT email, vote FROM votes WHERE track_id = ? ORDER BY ts", (self["id"], )):
            result[email] = vote
        return result

    @classmethod
    def from_file(cls, filename, labels=None):
        """Creates a new track from file, saves it.  If no labels were
        specified, adds the default ones."""
        filepath = os.path.join(ardj.settings.get_music_dir(), filename)
        tags = ardj.tags.get(filepath)

        t = cls(filename=filename)
        if "length" not in tags:
            raise Exception("No length in %s" % filename)
        t["artist"] = tags.get("artist", "Unknown Artist")
        t["title"] = tags.get(
            "title", os.path.basename(filename))
        t["length"] = tags["length"]
        t["weight"] = 1.0
        t["real_weight"] = 1.0
        t.put()

        if labels is None:
            labels = ardj.settings.get("default_labels")
        if labels:
            t.set_labels(labels)

        return t

    @classmethod
    def find_all_playlists(cls, count):
        sql = "SELECT %s FROM %s WHERE weight > 0 ORDER BY artist, title" % (
            cls._fields_sql(), cls.table_name)
        params = []
        return cls._fetch_rows(sql, params)

    @classmethod
    def query(cls, playlist=None, artist=None, tag=None, count=100):
        sql = "SELECT %s FROM %s WHERE weight > 0" % (
            cls._fields_sql(), cls.table_name)
        params = []
        order = "artist, title"

        if not playlist or playlist == "all":
            pass
        elif playlist == "never":
            sql += " AND last_played IS NULL"
        elif playlist == "recent":
            ts = time.time() - RECENT_SECONDS
            sql += " AND last_played > ?"
            params.append(int(ts))
            order = "last_played DESC"
        else:
            sql += " AND id IN (SELECT track_id FROM labels WHERE label = ?)"
            params.append(playlist)

        if artist and artist != "All artists":
            sql += " AND artist = ?"
            params.append(artist)

        if tag and tag != "All tags":
            sql += " AND id IN (SELECT track_id FROM labels WHERE label = ?)"
            params.append(tag)

        sql += " ORDER BY %s" % order

        if count:
            sql += " LIMIT %u" % count

        return cls._fetch_rows(sql, params)

class Queue(Model):
    """Stores information about a track to play out of regular order."""
    table_name = "queue"
    fields = "id", "track_id", "owner"
    key_name = "id"

class Token(Model):
    table_name = "tokens"
    fields = "token", "login", "login_type", "active"
    key_name = "token"

    @classmethod
    def create(cls, email, active=0):
        while True:
            tmp = random.randrange(111111, 999999)
            rows = fetchone("SELECT 1 FROM tokens WHERE token = ?", (tmp, ))
            if rows is not None:
                continue

            execute(
                "INSERT INTO tokens (token, login, login_type, active) VALUES (?, ?, 'email', ?)",
                (tmp,
                 email,
                 active))

            return cls(token=tmp,
                       login=email,
                       login_type="email",
                       active=active)

class Label(Model):
    table_name = "labels"
    fields = ("track_id", "email", "label")

    @classmethod
    def delete_by_name(cls, name):
        execute(
            "DELETE FROM `%s` WHERE `label` = ?" %
            cls.table_name, (name, ))

    @classmethod
    def find_all_names(cls):
        rows = fetch(
            "SELECT DISTINCT label FROM labels WHERE track_id IN (SELECT id FROM tracks WHERE weight > 0) ORDER BY label")
        return [r[0] for r in rows]

    @classmethod
    def find_never_played_names(cls):
        rows = fetch(
            "SELECT DISTINCT label FROM labels WHERE track_id IN (SELECT id FROM tracks WHERE weight > 0 AND last_played IS NULL) ORDER BY label")
        return [r[0] for r in rows]

    @classmethod
    def find_recently_played_names(cls):
        limit = time.time() - RECENT_SECONDS
        rows = fetch(
            "SELECT DISTINCT label FROM labels WHERE track_id IN (SELECT id FROM tracks WHERE weight > 0 AND last_played > ?) ORDER BY label",
            (limit,
             ))
        return [r[0] for r in rows]

    @classmethod
    def find_crossing_names(cls, name):
        rows = fetch(
            "SELECT DISTINCT label FROM labels WHERE track_id IN (SELECT id FROM tracks INNER JOIN labels l ON l.track_id = tracks.id WHERE weight > 0 AND l.label = ?)",
            (name,
             ))
        return [r[0] for r in rows]

    @classmethod
    def query_names(cls, playlist=None, artist=None):
        sql = "SELECT DISTINCT label FROM labels WHERE 1"
        params = []

        if playlist == "all":
            pass
        elif playlist == "never":
            sql += " AND track_id IN (SELECT id FROM tracks WHERE weight > 0 AND last_played IS NULL)"
        elif playlist == "recent":
            ts = time.time() - RECENT_SECONDS
            sql += " AND track_id IN  (SELECT id FROM tracks WHERE weight > 0 AND last_played > ?)"
            params.append(int(ts))
        else:
            sql += " AND track_id IN (SELECT id FROM tracks INNER JOIN labels l ON l.track_id = tracks.id WHERE weight > 0 AND l.label = ?)"
            params.append(playlist)

        if not artist or artist == "All artists":
            pass
        else:
            sql += " AND track_id IN (SELECT id FROM tracks WHERE weight > 0 AND artist = ?)"
            params.append(artist)

        sql += " ORDER BY label"

        return [r[0] for r in fetch(sql, params)]

class Artist(Model):
    """A virtual model without a table."""
    fields = ("name", )

    @classmethod
    def query(cls, playlist=None, tag=None):
        sql = "SELECT DISTINCT artist FROM tracks WHERE weight > 0"
        params = []

        if playlist == "all":
            pass
        elif playlist == "never":
            sql += " AND last_played IS NULL"
        elif playlist == "recent":
            sql += " AND last_played >= ?"
            params.append(time.time() - RECENT_SECONDS)
        else:
            sql += " AND id IN (SELECT track_id FROM labels WHERE label = ?)"
            params.append(playlist)

        if tag and tag != "All tags":
            sql += " AND id IN (SELECT track_id FROM labels WHERE label = ?)"
            params.append(tag)

        sql += " ORDER BY artist"
        return cls._fetch_rows(sql, params)

    @classmethod
    def query_all(cls):
        sql = "SELECT DISTINCT artist FROM tracks WHERE weight > 0 ORDER BY artist"
        return cls._fetch_rows(sql, ())

    @classmethod
    def find_unplayed(cls):
        sql = "SELECT DISTINCT artist FROM tracks WHERE weight > 0 AND last_played IS NULL ORDER BY artist"
        return cls._fetch_rows(sql, ())

    @classmethod
    def find_recent(cls):
        limit = time.time() - RECENT_SECONDS
        sql = "SELECT DISTINCT artist FROM tracks WHERE weight > 0 AND last_played > ? ORDER BY artist"
        return cls._fetch_rows(sql, (limit, ))

    @classmethod
    def find_by_tag(cls, tag):
        sql = "SELECT DISTINCT artist FROM tracks WHERE weight > 0 AND id IN (SELECT track_id FROM labels WHERE label = ?) ORDER BY artist"
        return cls._fetch_rows(sql, (tag, ))

class database:
    """
    Interface to the database.
    """
    instance = None

    def __init__(self, filename):
        """
        Opens the database, creates tables if necessary.
        """
        self.filename = filename
        isnew = not os.path.exists(self.filename)
        try:
            self.db = sqlite.connect(self.filename, check_same_thread=False)
        except Exception as e:
            logging.error('Could not open database %s: %s' % (filename, e))
            raise

        self.db.create_collation('UNICODE', ardj.util.ucmp)
        self.db.create_function('ULIKE', 2, self.sqlite_ulike)

    def __del__(self):
        self.commit()
        logging.debug('Database closed.')

    @classmethod
    def get_instance(cls):
        if cls.instance is None:
            filename = ardj.settings.getpath2(
                "database_path", "database/local")
            if filename is None:
                raise Exception('Config option database/local not set.')
            cls.instance = cls(filename)
        return cls.instance

    def sqlite_ulike(self, a, b):
        if a is None or b is None:
            return None
        if ardj.util.lower(b) in ardj.util.lower(a):
            return 1
        return 0

    def cursor(self):
        """Returns a new SQLite cursor, for internal use."""
        stack = "".join(traceback.format_stack()[-5:-1])
        #logging.debug("Starting a transaction (returning a cursor)\n%s" % stack)
        return self.db.cursor()

    def commit(self):
        """Commits current transaction, for internal use. """
        self.db.commit()

    def rollback(self):
        """Cancel pending changes."""
        logging.debug("Rolling back a transaction.")
        self.db.rollback()

    def fetch(self, sql, params=None):
        return self.execute(sql, params, fetch=True)

    def execute(self, sql, params=None, fetch=False):
        cur = self.db.cursor()
        try:
            args = [sql]
            if params is not None:
                args.append(params)
            cur.execute(*args)
            if fetch:
                return cur.fetchall()
            elif sql.startswith("INSERT "):
                return cur.lastrowid
            else:
                return cur.rowcount
        except BaseException:
            logging.exception(f"Failed SQL statement: {sql}, params: {params}")
            raise
        finally:
            cur.close()

    def update(self, table, args):
        """Performs update on a label.

        Updates the table with values from the args dictionary, key "id" must
        identify the record.  Example:

        db.update('tracks', {'weight': 1, 'id': 123})
        """
        sql = []
        params = []
        for k in args:
            if k != 'id':
                sql.append(k + ' = ?')
                params.append(args[k])
        params.append(args['id'])

        self.execute(
            'UPDATE %s SET %s WHERE id = ?' %
            (table, ', '.join(sql)), tuple(params))

    def debug(self, sql, params, quiet=False):
        """Logs the query in human readable form.

        Replaces question marks with parameter values (roughly)."""
        for param in params:
            param = str(param)
            if param.isdigit():
                sql = sql.replace('?', param, 1)
            else:
                sql = sql.replace('?', "'" + param + "'", 1)
        logging.debug(f"SQL: {sql}")
        return sql

    def purge(self):
        """Removes stale data.

        Stale data in queue items, labels and votes linked to tracks that no
        longer exist.  In addition to deleting such links, this function also
        analyzed all tables (to optimize indexes) and vacuums the database.
        """
        old_size = os.stat(self.filename).st_size
        self.execute(
            'DELETE FROM queue WHERE track_id NOT IN (SELECT id FROM tracks)')
        self.execute(
            'DELETE FROM labels WHERE track_id NOT IN (SELECT id FROM tracks)')
        self.execute(
            'DELETE FROM votes WHERE track_id NOT IN (SELECT id FROM tracks)')
        for table in ('playlists', 'tracks', 'queue',
                      'urgent_playlists', 'labels', 'karma'):
            self.execute('ANALYZE ' + table)
        logging.info('%u bytes saved after database purge.' %
                     (os.stat(self.filename).st_size - old_size))

    def mark_hitlist(self):
        """Marks best tracks with the "hitlist" label.

        Only processes tracks labelled with "music".  Before applying the
        label, removes it from the tracks that have it already."""
        set_label, check_label = 'hitlist', 'music'

        self.execute('DELETE FROM labels WHERE label = ?', (set_label, ))

        weight = fetchone(
            'SELECT real_weight FROM tracks WHERE id IN (SELECT track_id FROM labels WHERE label = ?) ORDER BY real_weight DESC LIMIT 19, 1',
            (check_label,
             ))
        if weight:
            self.execute(
                'INSERT INTO labels (track_id, label, email) SELECT id, ?, ? FROM tracks WHERE real_weight >= ? AND id IN (SELECT track_id FROM labels WHERE label = ?)',
                (set_label,
                 'ardj',
                 weight[0],
                    check_label,
                 ))

            lastfm = ardj.scrobbler.LastFM()
            lastfm.authorize()

            for artist, title in self.fetch(
                    'SELECT t.artist, t.title FROM tracks t INNER JOIN labels l ON l.track_id = t.id WHERE l.label = ?', (set_label, )):
                lastfm.love(artist, title)

    def mark_recent_music(self):
        """Marks last 100 tracks with "recent"."""
        self.execute('DELETE FROM labels WHERE label = ?', ('recent', ))
        self.execute(
            'INSERT INTO labels (track_id, label, email) SELECT id, ?, ? FROM tracks WHERE id IN (SELECT track_id FROM labels WHERE label = ?) ORDER BY id DESC LIMIT 100',
            ('recent',
             'ardj',
             'music',
             ))

        self.execute('DELETE FROM labels WHERE label = ?', ('fresh', ))
        self.execute(
            'INSERT INTO labels (track_id, label, email) SELECT id, ?, ? FROM tracks WHERE count < 10 AND weight > 0 AND id IN (SELECT track_id FROM labels WHERE label = ?)',
            ('fresh',
             'ardj',
             'music',
             ))

        count = self.execute(
            "SELECT COUNT(*) FROM labels WHERE label = ?", ('fresh', ))
        print('Found %u fresh songs.' % count)

    def mark_orphans(self, set_label='orphan', quiet=False):
        """Labels orphan tracks with "orphan".

        Orphans are tracks that don't belong to a playlist."""
        used_labels = []
        for playlist in ardj.settings.load().get_playlists():
            if 'labels' in playlist:
                labels = playlist['labels']
            else:
                labels = [playlist['name']]
            used_labels += [l for l in labels if not l.startswith('-')]
        used_labels = list(set(used_labels))

        if not(used_labels):
            logging.warning(
                'Could not mark orphan tracks: no labels are used in playlists.yaml')
            return False

        self.execute('DELETE FROM labels WHERE label = ?', (set_label, ))

        sql = 'SELECT id, artist, title FROM tracks WHERE weight > 0 AND id NOT IN (SELECT track_id FROM labels WHERE label IN (%s)) ORDER BY artist, title' % ', '.join([
            '?'] * len(used_labels))
        rows = self.fetch(sql, used_labels)

        if rows:
            if not quiet:
                print('%u orphan tracks found:' % len(rows))
            for row in rows:
                if not quiet:
                    print(
                        '%8u; %s -- %s' %
                        (row[0],
                         (row[1] or 'unknown').encode('utf-8'),
                            (row[2] or 'unknown').encode('utf-8')))
                self.execute(
                    'INSERT INTO labels (track_id, email, label) VALUES (?, ?, ?)', (int(
                        row[0]), 'ardj', set_label))
            return True

    def get_artist_names(self, label=None, weight=0):
        if label is None:
            rows = self.fetch(
                "SELECT DISTINCT artist FROM tracks WHERE weight > ?", (weight, ))
        else:
            rows = self.fetch(
                "SELECT DISTINCT artist FROM tracks WHERE weight > ? AND id IN (SELECT track_id FROM labels WHERE label = ?)",
                (weight,
                 label,
                 ))
        return [r[0] for r in rows]

def Open(filename=None):
    """Returns the active database instance."""
    return database.get_instance()

def commit():
    # ts = time.time()
    # logging.debug("Commit.")
    Open().commit()
    # logging.debug("Commit took %s seconds." % (time.time() - ts))

def rollback():
    Open().rollback()

def fetch(sql, params=None):
    return Open().fetch(sql, params)

def fetchone(sql, params=None):
    result = fetch(sql, params)
    if result:
        return result[0]

def fetchcol(sql, params=None):
    """Returns a list of first column values."""
    rows = fetch(sql, params)
    if rows:
        return [row[0] for row in rows]

def execute(*args, **kwargs):
    return Open().execute(*args, **kwargs)

def init_database():
    logging.info("Checking database integrity...")
    db = Open()
    cur = db.cursor()
    for statement in SQL_INIT:
        try:
            cur.execute(statement)
        except BaseException:
            logging.error("Init statement failed: %s" % statement)
            raise
    db.commit()

def cmd_console(*args):
    """Open database console (sqlite3)"""
    from subprocess import Popen
    from ardj import database
    init_database()
    Popen(["sqlite3", "-header", database.Open().filename]).wait()

def cmd_dedup_tracks():
    """Merge duplicate tracks"""
    from .tracks import dedup_by_filename
    count = dedup_by_filename(verbose=True)
    if count:
        print("Removed %u duplicate tracks." % count)
        commit()
    else:
        print("No duplicate tracks found.")

def cmd_init():
    """
    Initialize the database

    Initializes the configured database by executing a set of preconfigured SQL
    instructions.  This is non-destructive.  You should run this after you
    empty the database.
    """
    init_database()

def cmd_purge():
    """Delete stale data"""
    Open().purge()
    commit()

def cmd_stats():
    """Show database statistics"""
    tracks = Track.find_all()
    count = len(tracks)
    length = sum([t.get("length", 0) for t in tracks])
    print("%u tracks, %.1f hours." % (count, length / 60 / 60))

def cmd_vote_stats():
    """Show vote statistics"""
    daily = {}
    hourly = {}

    for vote in Vote.find_all():
        ts = time.gmtime(vote["ts"])

        if prefix is not None:
            if not time.strftime("%Y-%m-%d %H:%M:%S", ts).startswith(prefix):
                continue

        day = int(time.strftime("%w", ts)) or 7
        daily[day] = daily.get(day, 0) + 1

        hour = int(time.strftime("%H", ts))
        hourly[hour] = hourly.get(hour, 0) + 1

    def dump_votes(votes, prefix):
        total = float(sum(votes.values()))
        for k, v in list(votes.items()):
            print("%s,%u,%u" % (prefix, k, int(v)))

    dump_votes(daily, "D")
    dump_votes(hourly, "H")


#!/usr/bin/env python

from lucky.lang.build_ast import ASTBuilder
from lucky.lang.parser import Parser

class Runtime:
    def __init__(self, file: str, mode: int = 0):
        self.parser = Parser('grammar/main.lark')
        self.ast_builder = ASTBuilder()

        tree = self.parser.parse(file)
        import rich

        rich.print(tree)

        transformed = ASTBuilder().transform(tree)
        rich.print(transformed)


#coding:utf-8

import numpy as np
import Image
import imglib
import math

'''
Given a folder, reads all the JPG files as Images, converts each of
them in row vectors and returns a list with all the row vectors
'''
def create_image_database(folder,n):
	k = 1
	X = None
	while k<=n:
		img = Image.open("%s/img_%d.jpg" % (folder,k))
		if X is None:
			X = np.array( [imglib.img_to_rvector(img)] )
		else:
			X = np.append(X,[imglib.img_to_rvector(img)],axis=0)
		k+=1

	return X.T

'''
Read all the images following the original file hierarchy
from the AT&T Database
	att_faces
		s1
			1.pgm
			2.pgm
			...
		s2
			1.pgm
			2.pgm
			...
p: number of subjects
n: number of images per subject
tl: list of image indexes. These ones will not be use as the database,
	but as the test database
'''
def create_image_database_2(folder,p,n,tl):
	X = None
	Y = None
	for pi in range(1,p+1):
		k = 1
		while k<=n:
			img = Image.open("%s/s%d/%d.pgm" % (folder,pi,k))

			if k in tl:
				#Used as test Image
				if Y is None:
					Y = np.array( [imglib.img_to_rvector(img)] )
				else:
					Y = np.append(Y,[imglib.img_to_rvector(img)],axis=0)
			else:
				#Used as Database Image
				if X is None:
					X = np.array( [imglib.img_to_rvector(img)] )
				else:
					X = np.append(X,[imglib.img_to_rvector(img)],axis=0)
			k+=1

	return X.T,Y.T

'''
	n: Number of image per person
	p_db: List of people indexes presented at the database
	p_uk: List of people indexes will not be at the database
	tl: List of image indexes (for all database people) that it will be used
		as test images

	Return
	X: Database image column vectors
	rtcl: Recognition Test Cases List
'''
def create_image_database_3(folder,n,p_db,p_uk,tl):
	X = None
	Y = None
	metadb = []
	rtcl = []

	#Database People
	for p in p_db:
		k=1
		while k<=n:
			f_str = "%s/s%d/%d.pgm" % (folder,p,k)
			img = Image.open(f_str)
			if k in tl:
				#Used as test Image
				vc = np.array( [imglib.img_to_rvector(img)] )
				rtcl.append( create_rtc(ti=f_str,po=p,ii=k,uk=False) )

				if Y is None:
					Y = np.array( [imglib.img_to_rvector(img)] )
				else:
					Y = np.append(Y,[imglib.img_to_rvector(img)],axis=0)
			else:
				#Used as Database Image
				metadb.append( create_metadb(im=f_str,po=p,ii=k) )
				img_v = imglib.img_to_rvector(img)
				if X is None:
					X = np.array( [img_v] )
				else:
					X = np.append(X,[img_v],axis=0)

			k+=1

	db_images = len(p_db)*( n-len(tl) )
	test_images = len(p_uk)*n + len(p_db)*( len(tl) )

	#Unknown People
	for p in p_uk:
		for k in range(1,11):
			f_str = "%s/s%d/%d.pgm" % (folder,p,k)
			img = Image.open(f_str)

			vc = np.array( [imglib.img_to_rvector(img)] )
			rtcl.append( create_rtc(ti=f_str,po=p,ii=k,uk=True) )

			if Y is None:
				Y = np.array( [imglib.img_to_rvector(img)] )
			else:
				Y = np.append(Y,[imglib.img_to_rvector(img)],axis=0)

	return X.T,Y.T,metadb,rtcl

def create_rtc(ti=None,mi=None,po=None,pi=None,ii=None,di=None,uk=None,vc=None):
	return {"test_image":ti,"match_image":mi,"person_owner":po,"person_identified":pi,"image_index":ii,"distance":di,"unknown":uk,"vector":vc}

def create_metadb(im=None,po=None,ii=None):
	return {"image_name":im,"person_owner":po,"image_index":ii}


import numpy as np
from skimage import io
import matplotlib.pyplot as plt

from util import dict_minus_immutable, load_dict, timing
from config import *

@timing
def get_face_matches(known_faces, face):
    return np.linalg.norm(known_faces - face, axis=1)

@timing
def calculate_quality(own_matches, other_matches):
    """
    IDEAS:
        - ? drop the 10% worst images of ourselves
            > or probably not - we want to score high on these too
        - ! drop the 80% worst images of other
            > we are primarily interested in other images that look a bit like us
        - subtract standard deviation
    """
    details = {}
    avg_own = np.average(own_matches)
    avg_oth = np.average(other_matches)
    std_own = np.std(own_matches)
    std_oth = np.std(other_matches)

    details["avg_own"] = avg_own
    details["avg_oth"] = avg_oth
    details["std_own"] = std_own
    details["std_oth"] = std_oth

    others_like_me_n = int(len(other_matches)*0.20) # the 20% that looks most like me
    details["avg_others_like_me"] = np.average(sorted(other_matches)[:others_like_me_n])

    # Number of own images that fails comparison
    details["n_false_negative"] = len(own_matches[own_matches > 0.6])

    # Number of other images that tests positive
    details["n_false_positive"] = len(other_matches[other_matches < 0.6])

    quality = avg_oth - avg_own
    return (quality, details)

def find_other_embeddings(embeddings_per_person, person_id):
    others_embeddings = dict_minus_immutable(embeddings_per_person, person_id).values()
    return [item for sublist in others_embeddings for item in sublist]  # just efficient flatten()

def calculate_qualities(people):
    print("Calculating embedding qualities")
    embeddings_per_person = dict([(k, [f["embedding"] for f in v["faces"]]) for k, v in people.items()])

    for person_id, person in people.items():
        others_embeddings = find_other_embeddings(embeddings_per_person, person_id)
        print("  " + person_id + ": faces: ", len(person["faces"]), ", other faces: ", len(others_embeddings))

        for face in person["faces"]:
            own_embeddings_except_this_face = [f["embedding"] for f in person["faces"] if f["image_id"] != face["image_id"]]
            some_others_embeddings = others_embeddings # For now use all, it runs fast :)

            own_matches = get_face_matches(own_embeddings_except_this_face, face["embedding"])
            other_matches = get_face_matches(some_others_embeddings, face["embedding"])
            quality, quality_details = calculate_quality(own_matches, other_matches)

            face["quality"] = quality
            face["quality_details"] = quality_details

        # print("   - qualities: ", [face["quality"] for face in person["faces"]])

def filter_dataset_errors(people):
    """
    The MS-Celeb dataset have lots of people placed in wrong
    """
    calculate_qualities(people)     # Or use a tailored calculation for filtering???
    print("Filtering dataset for people not belonging")
    for person_id, person in people.items():
        faces = sorted(person["faces"], key=lambda f: -f['quality'])
        n_false_negatives = faces[0]["quality_details"]["n_false_negative"]
        if n_false_negatives > 0:
            print("Removing ",n_false_negatives," false negatives (probably wrong person) from ",person_id)
            person["faces"] = faces[:-n_false_negatives]

def print_best_image(people):
    for person_id, person in people.items():
        faces = sorted(person["faces"], key=lambda f: -f['quality'])
        print("Best image for recognition for "+person_id+" is "+faces[0]["image_id"]+"; having quality:", faces[0]["quality"])

def decorate_image_with_metrics(ax, image, face):
    image_height, image_width, unused = image.shape
    h, w = image_height / 5, image_width / 2   # text should be unaffected by image size (that may be uneven)
    # ax.text(0, 0, "{0:.2f}".format(face["quality"]), fontsize=10, color='red', verticalalignment='bottom')
    ax.text(0, 0,   "{0:.2f}".format(face["quality_details"]["avg_own"]), color='red', verticalalignment='top')
    ax.text(w, 0,   "{0:.2f}".format(face["quality_details"]["avg_oth"]), color='green', verticalalignment='top')
    ax.text(0, h,   "{}/{}".format(face["quality_details"]["n_false_negative"], face["quality_details"]["n_false_positive"]), color='blue', verticalalignment='top')
    ax.text(w, h,   "{0:.2f}".format(face["quality_details"]["avg_others_like_me"]), color='white', verticalalignment='top')

def plot_persons_faces(person):
    faces = sorted(person["faces"], key=lambda f: -f['quality'])    # best quality first
    faces = faces[0:18] + faces[-18:]   # Just some, not all
    num_cols = int(np.sqrt(len(faces))-(1e-7)) + 1
    print("Plotting face for ", person["person_id"])
    plt.figure(person["person_id"], figsize=(num_cols, num_cols))
    for i, face in enumerate(faces):
        image = io.imread(face["path"])
        ax = plt.subplot(num_cols, num_cols, i + 1)
        ax.set_axis_off()
        decorate_image_with_metrics(ax, image, face)
        plt.imshow(image)

    plt.subplots_adjust(wspace=0.1, hspace=0, left=0, right=1, top=1, bottom=0)
    plt.draw()  # plt.show()
    return None

def plot_images(people):
    # for person in people.values():
    #     plot_persons_faces(person)
    people_arr = list(people.values())
    plot_persons_faces(people_arr[1])
    plot_persons_faces(people_arr[5])
    plot_persons_faces(people_arr[11])
    plot_persons_faces(people_arr[12])
    plot_persons_faces(people_arr[13])

people = load_dict(intermediate_file)
#filter_dataset_errors(people)
calculate_qualities(people)

print_best_image(people)

plot_images(people)
plt.show()

# people:
#{
#    'm.03g19n': {
#        'person_id': 'm.03g19n',
#        'best_landmark':, ''
#        'faces': [
#            {
#                'image_id': '82-FaceId-0.jpg',
#                'quality': TBD: 0..1
#                'quality_details': {
#                   avg_own, avg_oth, std_own, std_oth, ...
#                },
#                'path':  '/Users/trygve/data/ms-celeb/MsCelebV1-Faces-Aligned.Samples/samples//m.03g19n/82-FaceId-0.jpg',
#                'bounds': rectangle(21,37,93,109),
#                'landmarks-5': [(80, 52), (66, 54), (31, 52), (44, 54), (56, 84)],
#                'landmarks-68': [(80, 52), (66, 54), (31, 52), (44, 54), (56, 84), ...],
#                'embedding': array([<128-vector embedding>])
#            },
#            {
#                'image_id': '91-FaceId-0.jpg',
#                'quality': TBD: 0..1
#                'quality_details': {
#                   avg_own, avg_oth, std_own, std_oth, ...
#                },
#                'path':  '/Users/trygve/data/ms-celeb/MsCelebV1-Faces-Aligned.Samples/samples//m.03g19n/91-FaceId-0.jpg',
#                'bounds': rectangle(21,37,93,109),
#                'landmarks-5': [(80, 52), (66, 54), (31, 52), (44, 54), (56, 84)],
#                'landmarks-68': [(80, 52), (66, 54), (31, 52), (44, 54), (56, 84), ...],
#                'embedding': array([<128-vector embedding>])
#            },
#            ...
#        ]
#    },
#    'm.044wfaf': {...}
#}


from django.contrib import admin
from django.urls import path,include
from . import views

urlpatterns = [
    path("Home/",views.IndexPage,name="homepage"),
    path("Register-Page/",views.RegisterPage,name="regpage"),
    path("Registration/",views.Register,name="Registerdata"),
    path("Updaate-Page/<int:pk>",views.UpdatePage,name="updatepage"),
    path("Updaate/<int:pk>",views.Update,name="update"),
    path("Delete/<int:pk>",views.Delete,name="delete"),
    path('',views.LoginPage,name='loginpage'),
    path('Login-User/',views.LoginUser,name='loginuser'),
    path('Logout/',views.LogOut,name='logout'),
]


from model import ModelBase

class EmailThread(ModelBase):
    def __init__(self, id_, account_id):
        self.id_ = id_
        self.account_id = account_id
        self.emails = list()

    def compute_score(self):
        email = self.emails[0]
        if self.has_newer(email.sent_at.to_f()):
            self.logging.info("Keeping old score for " + self.log_ident )
            return False
        else:
            self.score = email.score
            return True

    def has_newer(self, sent_at):
        raise "Must override has_newer"


import tkinter as tk
from tkinter import messagebox
from operator import attrgetter

class HighScore:
    def __init__(self, name, size, time):
        self.name = name
        self.size = size
        self.time = time

def read_highscores_from_file(file_name):
    highscores = []
    try:
        with open(file_name, 'r') as file:
            for line in file:
                data = line.strip().split(',')
                if len(data) == 3:
                    name, size, time = data
                    rows, cols = map(int, size.split('x'))
                    highscores.append(HighScore(name, (rows, cols), float(time)))
    except FileNotFoundError:
        print(f"File '{file_name}' not found.")
    return highscores

def write_highscores_to_file(file_name, highscores):
    try:
        with open(file_name, 'w') as file:
            for highscore in highscores:
                size = f"{highscore.size[0]}x{highscore.size[1]}"
                file.write(f"{highscore.name},{size},{highscore.time}\n")
    except IOError:
        print(f"Error writing to '{file_name}'.")

def sort_highscores(scores):
    scores.sort(key=lambda x: (-x.size[0], -x.size[1], x.time))

def show_highscores():
    sorted_scores = read_highscores_from_file('highscores.txt')
    sort_highscores(sorted_scores)

    top = tk.Toplevel(root)
    top.title("High Scores")

    for i, score in enumerate(sorted_scores, start=1):
        label = tk.Label(top, text=f"{i}. Player: {score.name}, Size: {score.size[0]}x{score.size[1]}, Time: {score.time}")
        label.pack()

root = tk.Tk()
root.title("High Scores")

button = tk.Button(root, text="Show High Scores", command=show_highscores)
button.pack()

root.mainloop()


def increment_string(strng):
    numstr=''
    strstr=''
    str2=strstr[:len(numstr)]
    print(str2)
    for cha in strng:
        if cha.isdigit():
            numstr+=cha
        else:
            strstr+=cha
    if numstr=='':
        print(strstr+'1')
        return 1
    num=int(numstr)
    str2=strstr[:len(numstr)+1]
    print(str2)
    # print(numstr)

    if len(str(num))==len(numstr):
        print(strstr+str(num+1))
    else:
        zeros=len(numstr)-len(str(num+1))
        print(zeros)
        res=('0'*zeros)+str(num+1)
        print(strstr+res)
        return strng

increment_string("foo09")


from brownie import PayMeEthereum
from scripts.helper import get_account

def fund():
    pay_me = PayMeEthereum[-1]
    account = get_account()
    entrance_fee = pay_me.getEntranceFee()
    print(entrance_fee)
    print(f"The current entry fee is {entrance_fee}")
    print("Funding")
    pay_me.buyMeACoffee({"from": account, "value": entrance_fee})

def withdraw():
    pay_me = PayMeEthereum[-1]
    account = get_account()
    pay_me.withdraw({"from": account})

def main():
    fund()
    withdraw()


import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from sklearn.cluster import KMeans
from sklearn.cluster import MeanShift, estimate_bandwidth
from itertools import cycle

data=pd.read_csv("table_jones.csv", delimiter=',', header=None)
x=data[1].values
y=data[4].values
XX=np.array(list(zip(x,y)))

# #############################################################################
# Compute clustering with MeanShift

# The following bandwidth can be automatically detected using
bandwidth = estimate_bandwidth(XX, quantile=0.2, n_samples=500)

ms = MeanShift(bandwidth=bandwidth, bin_seeding=True)
ms.fit(XX)
labels = ms.labels_
cluster_centers = ms.cluster_centers_

labels_unique = np.unique(labels)
n_clusters_ = len(labels_unique)

print("number of estimated clusters : %d" % n_clusters_)

# #############################################################################
colors = cycle('bgrcmykbgrcmykbgrcmykbgrcmyk')
for k, col in zip(range(n_clusters_), colors):
    my_members = labels == k
    cluster_center = cluster_centers[k]
    plt.plot(XX[my_members, 0], XX[my_members, 1], col + '.')
    plt.plot(cluster_center[0], cluster_center[1], 'o', markerfacecolor=col,
             markeredgecolor='k', markersize=14)

plt.xlabel('Absolute magnitude')
plt.ylabel('Median age [Myr]')
plt.title('(3152) Jones')
plt.text(15.2,4.65,'(a)', fontsize=12)
plt.savefig('/home/valerio/PRAVEC/TEXT/ML_JONES/MEAN-SHIFT_JONES/Mean-shift_jones.eps')
plt.show()


#!/usr/bin/env python3

""" The Bank model data
"""

import enum

from financial_game.table import Table, Identifier, String, Enum, ForeignKey
from financial_game.database import Join

class TypeOfBank(enum.Enum):
    """Types of bank objects"""

    BANK = 1  # bank

class Bank(Table):
    """bank info"""

    _db = None
    id = Identifier()
    name = String(50, allow_null=False)
    type = Enum(TypeOfBank, allow_null=False)
    url = String(2083)

    def __eq__(self, other):
        assert isinstance(other, Bank)
        return other.id == self.id

    @staticmethod
    def create(name: str, url: str = None, bank_type=TypeOfBank.BANK):
        """Create a new bank"""
        assert name is not None
        bank = Bank(name=name, url=url, type=bank_type, _normalize_=False).denormalize()
        created = Bank._db.insert(Table.name(Bank), **bank)
        return Bank(**created)

    @staticmethod
    def fetch(bank_id: int):
        """Get a bank by its id"""
        found = Bank._db.get_one_or_none(
            Table.name(Bank), _where_="id = :bank_id", bank_id=bank_id
        )
        return None if found is None else Bank(**found)

    @staticmethod
    def every(bank_type=TypeOfBank.BANK):
        """Get list of all banks"""
        return [
            Bank(**u)
            for u in Bank._db.get_all(
                Table.name(Bank), _where_="type = :bank_type", bank_type=bank_type.name
            )
        ]

    @staticmethod
    def total(bank_type=TypeOfBank.BANK):
        """Count total banks"""
        return Bank._db.get_one_or_none(
            Table.name(Bank),
            "COUNT(*)",
            _where_="type = :bank_type",
            bank_type=bank_type.name,
        )["COUNT(*)"]

    def account_types(self):
        """get the types of accounts at the bank"""
        found = Bank._db.get_all(
            Table.name(AccountType), _where_="bank_id = :bank_id", bank_id=self.id
        )
        return [AccountType(**t) for t in found]

class TypeOfAccount(enum.Enum):
    """Types of account_type objects"""

    CRED = 1  # Credit Card
    CHCK = 2  # Checking
    SAVE = 3  # Savings
    MONM = 4  # Money Market
    BROK = 5  # Brokrage account

class AccountType(Table):
    """Bank account brand"""

    _db = None
    id = Identifier()
    name = String(50, allow_null=False)
    type = Enum(TypeOfAccount, allow_null=False)
    url = String(2083)
    bank_id = ForeignKey(Bank, allow_null=False)

    def __eq__(self, other):
        assert isinstance(other, AccountType)
        return other.id == self.id

    @staticmethod
    def create(bank, name: str, account_type: TypeOfAccount, url: str = None):
        """Create a new bank"""
        assert name is not None
        assert isinstance(bank, (int, Bank))
        account_type = AccountType(
            name=name,
            url=url,
            bank_id=bank.id if isinstance(bank, Bank) else bank,
            type=account_type,
            _normalize_=False,
        ).denormalize()
        created = AccountType._db.insert(Table.name(AccountType), **account_type)
        return AccountType(**created)

    @staticmethod
    def every(bank_type=TypeOfBank.BANK):
        """Return ([Bank], {bank.id:[AccountType]})"""
        results = AccountType._db.get_all(
            Table.name(AccountType),
            f"{Table.name(Bank)}.id as bank_id",
            f"{Table.name(Bank)}.name as bank_name",
            f"{Table.name(Bank)}.type as bank_type",
            f"{Table.name(Bank)}.url as bank_url",
            f"{Table.name(AccountType)}.id as account_type_id",
            f"{Table.name(AccountType)}.name as account_type_name",
            f"{Table.name(AccountType)}.type as account_type_type",
            f"{Table.name(AccountType)}.url as account_type_url",
            _join_=Join(
                Table.name(Bank),
                f"{Table.name(Bank)}.id = {Table.name(AccountType)}.bank_id",
            ),
            _where_=f"{Table.name(Bank)}.type = :bank_type",
            bank_type=bank_type.name,
        )
        account_types = {}
        banks = []

        for account_type in results:
            bank = Bank(
                id=account_type["bank_id"],
                name=account_type["bank_name"],
                type=account_type["bank_type"],
                url=account_type["bank_url"],
            )

            if bank not in banks:
                banks.append(bank)

            account_types[bank.id] = account_types.get(bank.id, [])
            account_types[bank.id].append(
                AccountType(
                    id=account_type["account_type_id"],
                    name=account_type["account_type_name"],
                    type=account_type["account_type_type"],
                    url=account_type["account_type_url"],
                )
            )

        return banks, account_types

    @staticmethod
    def fetch(account_type_id: int):
        """Get an account type by its id"""
        found = AccountType._db.get_one_or_none(
            Table.name(AccountType),
            _where_="id = :account_type_id",
            account_type_id=account_type_id,
        )
        return None if found is None else AccountType(**found)

    def bank(self):
        """Get the bank this type is for"""
        found = AccountType._db.get_one_or_none(
            Table.name(Bank), _where_="id = :bank_id", bank_id=self.bank_id
        )
        return Bank(**found)


import argparse
import os

from matplotlib.gridspec import GridSpec
import matplotlib.pyplot as plt
import seaborn as sns

from jax import lax, random, vmap
from jax.config import config as jax_config
from jax.experimental import optimizers
import jax.numpy as np
from jax.tree_util import tree_map

from numpyro.contrib.autoguide import AutoIAFNormal
from numpyro.diagnostics import summary
import numpyro.distributions as dist
from numpyro.distributions.util import logsumexp
from numpyro.handlers import sample
from numpyro.hmc_util import initialize_model
from numpyro.mcmc import mcmc
from numpyro.svi import elbo, svi

# TODO: remove when the issue https://github.com/google/jax/issues/939 is fixed upstream
# The behaviour when training guide under fast math mode is unstable.
os.environ["XLA_FLAGS"] = "--xla_cpu_enable_fast_math=false"

"""
This example illustrates how to use a trained AutoIAFNormal autoguide to transform a posterior to a
Gaussian-like one. The transform will be used to get better mixing rate for NUTS sampler.

[1] Hoffman, M. et al. (2019), ["NeuTra-lizing Bad Geometry in Hamiltonian Monte Carlo Using Neural Transport"]
    (https://arxiv.org/abs/1903.03704).
"""

def dual_moon_pe(x):
    term1 = 0.5 * ((np.linalg.norm(x, axis=-1) - 2) / 0.4) ** 2
    term2 = -0.5 * ((x[..., :1] + np.array([-2., 2.])) / 0.6) ** 2
    return term1 - logsumexp(term2, axis=-1)

def dual_moon_model():
    x = sample('x', dist.Uniform(-4 * np.ones(2), 4 * np.ones(2)))
    pe = dual_moon_pe(x)
    sample('log_density', dist.Delta(log_density=-pe), obs=0.)

def make_transformed_pe(potential_fn, transform, unpack_fn):
    def transformed_potential_fn(z):
        u, intermediates = transform.call_with_intermediates(z)
        logdet = transform.log_abs_det_jacobian(z, u, intermediates=intermediates)
        return potential_fn(unpack_fn(u)) + logdet

    return transformed_potential_fn

def main(args):
    jax_config.update('jax_platform_name', args.device)

    print("Start vanilla HMC...")
    vanilla_samples = mcmc(args.num_warmup, args.num_samples, init_params=np.array([2., 0.]),
                           potential_fn=dual_moon_pe, progbar=True)

    opt_init, opt_update, get_params = optimizers.adam(0.001)
    rng_guide, rng_init, rng_train = random.split(random.PRNGKey(1), 3)
    guide = AutoIAFNormal(rng_guide, dual_moon_model, get_params, hidden_dims=[args.num_hidden])
    svi_init, svi_update, _ = svi(dual_moon_model, guide, elbo, opt_init, opt_update, get_params)
    opt_state, _ = svi_init(rng_init)

    def body_fn(val, i):
        opt_state_, rng_ = val
        loss, opt_state_, rng_ = svi_update(i, rng_, opt_state_)
        return (opt_state_, rng_), loss

    print("Start training guide...")
    (last_state, _), losses = lax.scan(body_fn, (opt_state, rng_train), np.arange(args.num_iters))
    print("Finish training guide. Extract samples...")
    guide_samples = guide.sample_posterior(random.PRNGKey(0), last_state,
                                           sample_shape=(args.num_samples,))

    transform = guide.get_transform(last_state)
    unpack_fn = guide.unpack_latent

    _, potential_fn, constrain_fn = initialize_model(random.PRNGKey(0), dual_moon_model)
    transformed_potential_fn = make_transformed_pe(potential_fn, transform, unpack_fn)
    transformed_constrain_fn = lambda x: constrain_fn(unpack_fn(transform(x)))  # noqa: E731

    init_params = np.zeros(guide.latent_size)
    print("\nStart NeuTra HMC...")
    zs = mcmc(args.num_warmup, args.num_samples, init_params, potential_fn=transformed_potential_fn)
    print("Transform samples into unwarped space...")
    samples = vmap(transformed_constrain_fn)(zs)
    summary(tree_map(lambda x: x[None, ...], samples))

    # make plots

    # IAF guide samples (for plotting)
    iaf_base_samples = dist.Normal(np.zeros(2), 1.).sample(random.PRNGKey(0), (1000,))
    iaf_trans_samples = vmap(transformed_constrain_fn)(iaf_base_samples)['x']

    x1 = np.linspace(-3, 3, 100)
    x2 = np.linspace(-3, 3, 100)
    X1, X2 = np.meshgrid(x1, x2)
    P = np.clip(np.exp(-dual_moon_pe(np.stack([X1, X2], axis=-1))), a_min=0.)

    fig = plt.figure(figsize=(12, 16), constrained_layout=True)
    gs = GridSpec(3, 2, figure=fig)
    ax1 = fig.add_subplot(gs[0, 0])
    ax2 = fig.add_subplot(gs[0, 1])
    ax3 = fig.add_subplot(gs[1, 0])
    ax4 = fig.add_subplot(gs[1, 1])
    ax5 = fig.add_subplot(gs[2, 0])
    ax6 = fig.add_subplot(gs[2, 1])

    ax1.plot(np.log(losses[1000:]))
    ax1.set_title('Autoguide training log loss (after 1000 steps)')

    ax2.contourf(X1, X2, P, cmap='OrRd')
    sns.kdeplot(guide_samples['x'][:, 0].copy(), guide_samples['x'][:, 1].copy(), n_levels=30, ax=ax2)
    ax2.set(xlim=[-3, 3], ylim=[-3, 3],
            xlabel='x0', ylabel='x1', title='Posterior using AutoIAFNormal guide')

    sns.scatterplot(iaf_base_samples[:, 0], iaf_base_samples[:, 1], ax=ax3, hue=iaf_trans_samples[:, 0] < 0.)
    ax3.set(xlim=[-3, 3], ylim=[-3, 3],
            xlabel='x0', ylabel='x1', title='AutoIAFNormal base samples (True=left moon; False=right moon)')

    ax4.contourf(X1, X2, P, cmap='OrRd')
    sns.kdeplot(vanilla_samples[:, 0].copy(), vanilla_samples[:, 1].copy(), n_levels=30, ax=ax4)
    ax4.plot(vanilla_samples[-50:, 0], vanilla_samples[-50:, 1], 'bo-', alpha=0.5)
    ax4.set(xlim=[-3, 3], ylim=[-3, 3],
            xlabel='x0', ylabel='x1', title='Posterior using vanilla HMC sampler')

    sns.scatterplot(zs[:, 0], zs[:, 1], ax=ax5, hue=samples['x'][:, 0] < 0.,
                    s=30, alpha=0.5, edgecolor="none")
    ax5.set(xlim=[-5, 5], ylim=[-5, 5],
            xlabel='x0', ylabel='x1', title='Samples from the warped posterior - p(z)')

    ax6.contourf(X1, X2, P, cmap='OrRd')
    sns.kdeplot(samples['x'][:, 0].copy(), samples['x'][:, 1].copy(), n_levels=30, ax=ax6)
    ax6.plot(samples['x'][-50:, 0], samples['x'][-50:, 1], 'bo-', alpha=0.2)
    ax6.set(xlim=[-3, 3], ylim=[-3, 3],
            xlabel='x0', ylabel='x1', title='Posterior using NeuTra HMC sampler')

    plt.savefig("neutra.pdf")
    plt.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="NeuTra HMC")
    parser.add_argument('-n', '--num-samples', nargs='?', default=20000, type=int)
    parser.add_argument('--num-warmup', nargs='?', default=0, type=int)
    parser.add_argument('--num-hidden', nargs='?', default=20, type=int)
    parser.add_argument('--num-iters', nargs='?', default=200000, type=int)
    parser.add_argument('--device', default='cpu', type=str, help='use "cpu" or "gpu".')
    args = parser.parse_args()
    main(args)