#!/usr/bin/env python3

'''Take a prefix and model name run predictions, and generate evaluations for crystal, bestonly,
and all test sets (take max affinity; if pose score is available also consider
max pose score).
Generates graphs and overall CV results.  Takes the prefix and (for now) assumes trial 0.
Will evaluate 100k model and best model prior to 100k, 50k and 25k
'''

import numpy as np
import os, sys
#os.environ["GLOG_minloglevel"] = "10"
sys.path.append("/home/dkoes/git/gninascripts/")
sys.path.append("/net/pulsar/home/koes/dkoes/git/gninascripts/")

import train, predict
import matplotlib, caffe
import matplotlib.pyplot as plt
import glob, re, sklearn, collections, argparse, sys
import sklearn.metrics
import scipy.stats

def evaluate_fold(testfile, caffemodel, modelname, datadir='../..',hasrmsd=False):
    '''Evaluate the passed model and the specified test set.
    Returns tuple:
    (correct, prediction, receptor, ligand, label (optional), posescore (optional))
    label and posescore are only provided is trained on pose data
    '''
    if not os.path.exists(modelname):
       print(modelname,"does not exist")

    caffe.set_mode_gpu()
    test_model = 'predict.%d.prototxt' % os.getpid()
    train.write_model_file(test_model, modelname, testfile, testfile, datadir)
    test_net = caffe.Net(test_model, caffemodel, caffe.TEST)
    lines = open(testfile).readlines()
    res = None
    i = 0 #index in batch
    correct = 0
    prediction = 0
    receptor = ''
    ligand = ''
    label = 0
    posescore = -1
    ret = []
    for line in lines:
        #check if we need a new batch of results
        if not res or i >= batch_size:
            res = test_net.forward()
            if 'output' in res:
                batch_size = res['output'].shape[0]
            else:
                batch_size = res['affout'].shape[0]
            i = 0

        if 'labelout' in res:
            label = float(res['labelout'][i])

        if 'output' in res:
            posescore = float(res['output'][i][1])

        if 'affout' in res:
            correct = float(res['affout'][i])

        if 'predaff' in res:
            prediction = float(res['predaff'][i])
            if not np.isfinite(prediction).all():
                os.remove(test_model)
                return [] #gracefully handle nan?

        #extract ligand/receptor for input file
        tokens = line.split()
        rmsd = -1
        for t in range(len(tokens)):
            if tokens[t].lower()=='none':
                #Flag that none as the receptor file, for ligand-only models
                ligand=tokens[t+1]

                #we assume that ligand is rec/<ligname>
                #set if correct, bail if not.
                m=re.search(r'(\S+)/(\S+)gninatypes',ligand)

                #Check that the match is not none, and that ligand ends in gninatypes
                if m is not None:
                    receptor=m.group(1)
                else:
                    print('Error: none receptor detected and ligand is improperly formatted.')
                    print('Ligand must be formatted: <rec>/<ligfile>.gninatypes')
                    print('Bailing.')
                    sys.exit(1)
                break

            elif tokens[t].endswith('gninatypes'):
                receptor = tokens[t]
                ligand = tokens[t+1]
                break
        if hasrmsd:
            rmsd = float(tokens[2])
        #(correct, prediction, receptor, ligand, label (optional), posescore (optional))
        if posescore < 0:
            ret.append((correct, prediction, receptor, ligand))
        elif hasrmsd:
            ret.append((correct, prediction, receptor, ligand, label, posescore, rmsd))
        else:
            ret.append((correct, prediction, receptor, ligand, label, posescore))

        i += 1 #batch index

    os.remove(test_model)
    return ret

def reduce_results(results, index):
    '''Return results with only one tuple for every receptor value,
    taking the one with the max value at index in the tuple (predicted affinity or pose score)
    '''
    res = dict() #indexed by receptor
    for r in results:
        name = r[2]
        if name not in res:
            res[name] = r
        elif res[name][index] < r[index]:
            res[name] = r
    return list(res.values())

def analyze_results(results, outname, uniquify=None):
    '''Compute error metrics from resuls.  RMSE, Pearson, Spearman.
    If uniquify is set, AUC and top-1 percentage are also computed,
    uniquify can be None, 'affinity', or 'pose' and is set with
    the all training set to select the pose used for scoring.
    Returns tuple:
    (RMSE, Pearson, Spearman, AUCpose, AUCaffinity, top-1)
    Writes (correct,prediction) pairs to outname.predictions
    '''

    #calc auc before reduction
    if uniquify and len(results[0]) > 5:
        labels = np.array([r[4] for r in results])
        posescores = np.array([r[5] for r in results])
        predictions = np.array([r[1] for r in results])
        aucpose = sklearn.metrics.roc_auc_score(labels, posescores)
        aucaff = sklearn.metrics.roc_auc_score(labels, predictions)

    if uniquify == 'affinity':
        results = reduce_results(results, 1)
    elif uniquify == 'pose':
        results = reduce_results(results, 5)

    predictions = np.array([r[1] for r in results])
    correctaff = np.array([abs(r[0]) for r in results])
    #(correct, prediction, receptor, ligand, label (optional), posescore (optional))

    rmse = np.sqrt(sklearn.metrics.mean_squared_error(correctaff, predictions))
    R = scipy.stats.pearsonr(correctaff, predictions)[0]
    S = scipy.stats.spearmanr(correctaff, predictions)[0]
    out = open('%s.predictions'%outname,'w')
    for (c,p) in zip(correctaff,predictions):
        out.write('%f %f\n' % (c,p))
    out.write('#RMSD %f\n'%rmse)
    out.write('#R %f\n'%R)

    if uniquify and len(results[0]) > 5:
        labels = np.array([r[4] for r in results])
        top = np.count_nonzero(labels > 0)/float(len(labels))
        return (rmse, R, S, aucpose, aucaff, top)
    else:
        return (rmse, R, S)

if __name__ == '__main__':
    if len(sys.argv) <= 4:
        print("Need caffemodel prefix,  modelname, output name and test prefixes (which should include _<slicenum>_ at end)")
        sys.exit(1)

    name = sys.argv[1]
    modelname = sys.argv[2]
    out = open(sys.argv[3],'w')

    allresults = []
    last = None
    #for each test dataset
    for testprefix in sys.argv[4:]:
        m = re.search('([^/ ]*)_(\d+)_$', testprefix)
        print(m,testprefix)
        if not m:
            print(testprefix,"does not end in slicenum")
        slicenum = int(m.group(2))
        testname = m.group(1)
        #find the relevant models for each fold
        testresults = {'best25': [], 'best50': [], 'best100': [], 'last': [], 'best250': [] }
        for fold in [0,1,2]:
            best25 = 0
            best50 = 0
            best100 = 0
            best250 = 0
            lastm = 0
            #identify best iteration models at each cut point for this fold
            for model in glob.glob('%s.%d_iter_*.caffemodel'%(name,fold)):
                m = re.search(r'_iter_(\d+).caffemodel', model)
                inum = int(m.group(1))
                if inum < 25000 and inum > best25:
                    best25 = inum
                if inum < 50000 and inum > best50:
                    best50 = inum
                if inum < 100000 and inum > best100:
                    best100 = inum
                if inum < 250000 and inum > best250:
                    best250 = inum
                if inum > lastm:
                    lastm = inum
            #evalute this fold
            testfile = '../types/%stest%d.types' % (testprefix,fold)
            #todo, avoid redundant repetitions
            if best25 > 0: testresults['best25'] += evaluate_fold(testfile, '%s.%d_iter_%d.caffemodel' % (name,fold,best25), modelname)
            if best50 > 0: testresults['best50'] += evaluate_fold(testfile, '%s.%d_iter_%d.caffemodel' % (name,fold,best50), modelname)
            if best100 > 0: testresults['best100'] += evaluate_fold(testfile, '%s.%d_iter_%d.caffemodel' % (name,fold,best100), modelname)
            if best250 > 0: testresults['best250'] += evaluate_fold(testfile, '%s.%d_iter_%d.caffemodel' % (name,fold,best250), modelname)
            if lastm > 0: testresults['last'] += evaluate_fold(testfile, '%s.%d_iter_%d.caffemodel' % (name,fold,lastm), modelname)

        for n in list(testresults.keys()):
            if len(testresults[n]) == 0:
                continue
            if len(testresults[n][0]) == 6:
                allresults.append( ('%s_pose'%testname, n) + analyze_results(testresults[n],('%s_pose_'%testname)+name+'_'+n,'pose'))
            allresults.append( ('%s_affinity'%testname, n) + analyze_results(testresults[n],('%s_affinity_'%testname)+name+'_'+n,'affinity'))

    for a in allresults:
        out.write(' '.join(map(str,a))+'\n')


# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import logging
from pathlib import Path

import librosa
import numpy as np
import paddle
import soundfile as sf
import yaml
from parakeet.data.get_feats import LogMelFBank
from parakeet.models.parallel_wavegan import PWGGenerator
from parakeet.models.parallel_wavegan import PWGInference
from parakeet.modules.normalizer import ZScore
from yacs.config import CfgNode as Configuration

def get_cfg_default():
    config_path = (Path(__file__).parent / "conf" / "default.yaml").resolve()
    with open(config_path, 'rt') as f:
        _C = yaml.safe_load(f)
        _C = Configuration(_C)
    config = _C.clone()
    return config

def evaluate(args, config):
    # dataloader has been too verbose
    logging.getLogger("DataLoader").disabled = True

    vocoder = PWGGenerator(**config["generator_params"])
    state_dict = paddle.load(args.checkpoint)
    vocoder.set_state_dict(state_dict["generator_params"])
    vocoder.remove_weight_norm()
    vocoder.eval()
    print("model done!")

    stat = np.load(args.stat)
    mu, std = stat
    mu = paddle.to_tensor(mu)
    std = paddle.to_tensor(std)
    normalizer = ZScore(mu, std)

    pwg_inference = PWGInference(normalizer, vocoder)

    input_dir = Path(args.input_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    mel_extractor = LogMelFBank(
        sr=config.fs,
        n_fft=config.n_fft,
        hop_length=config.n_shift,
        win_length=config.win_length,
        window=config.window,
        n_mels=config.n_mels,
        fmin=config.fmin,
        fmax=config.fmax)

    for utt_name in os.listdir(input_dir):
        wav, _ = librosa.load(str(input_dir / utt_name), sr=config.fs)
        # extract mel feats
        mel = mel_extractor.get_log_mel_fbank(wav)
        mel = paddle.to_tensor(mel)
        with paddle.no_grad():
            gen_wav = pwg_inference(mel)
        sf.write(
            str(output_dir / ("gen_" + utt_name)),
            gen_wav.numpy(),
            samplerate=config.fs)
        print(f"{utt_name} done!")

def main():
    # parse args and config and redirect to train_sp
    parser = argparse.ArgumentParser(
        description="Synthesize with parallel wavegan.")

    parser.add_argument(
        "--config", type=str, help="config file to overwrite default config.")
    parser.add_argument("--checkpoint", type=str, help="snapshot to load.")
    parser.add_argument(
        "--stat",
        type=str,
        help="mean and standard deviation used to normalize spectrogram when training parallel wavegan."
    )
    parser.add_argument("--input-dir", type=str, help="input dir of wavs.")
    parser.add_argument("--output-dir", type=str, help="output dir.")
    parser.add_argument(
        "--device", type=str, default="gpu", help="device to run.")
    parser.add_argument("--verbose", type=int, default=1, help="verbose.")

    args = parser.parse_args()

    paddle.set_device(args.device)

    config = get_cfg_default()
    if args.config:
        config.merge_from_file(args.config)

    print("========Args========")
    print(yaml.safe_dump(vars(args)))
    print("========Config========")
    print(config)

    evaluate(args, config)

if __name__ == "__main__":
    main()


import numpy as np
from keras import backend as K
from keras.engine.topology import Layer
from keras.layers.pooling import _GlobalPooling1D
from keras.layers import Conv1D, Input, LocallyConnected1D
from keras.layers.core import Dropout
from concise.utils.plot import seqlogo, seqlogo_fig
import matplotlib.pyplot as plt
from keras.engine import InputSpec

from concise.utils.pwm import DEFAULT_BASE_BACKGROUND, pssm_array2pwm_array, _pwm2pwm_info
from keras import activations
from keras import constraints
from concise import initializers
from concise import regularizers
from concise.regularizers import GAMRegularizer, SplineSmoother
from concise.utils.splines import BSpline
from concise.utils.helper import get_from_module
from concise.utils.plot import heatmap
from concise.preprocessing.sequence import (DNA, RNA, AMINO_ACIDS,
                                            CODONS, STOP_CODONS)
from concise.preprocessing.structure import RNAplfold_PROFILES

# --------------------------------------------
# Input()

def InputDNA(seq_length, name=None, **kwargs):
    """Input placeholder for array returned by `encodeDNA` or `encodeRNA`

    Wrapper for: `keras.layers.Input((seq_length, 4), name=name, **kwargs)`
    """
    return Input((seq_length, 4), name=name, **kwargs)

InputRNA = InputDNA

def InputCodon(seq_length, ignore_stop_codons=True, name=None, **kwargs):
    """Input placeholder for array returned by `encodeCodon`

    Note: The seq_length is divided by 3

    Wrapper for: `keras.layers.Input((seq_length / 3, 61 or 61), name=name, **kwargs)`
    """
    if ignore_stop_codons:
        vocab = CODONS
    else:
        vocab = CODONS + STOP_CODONS

    assert seq_length % 3 == 0
    return Input((seq_length / 3, len(vocab)), name=name, **kwargs)

def InputAA(seq_length, name=None, **kwargs):
    """Input placeholder for array returned by `encodeAA`

    Wrapper for: `keras.layers.Input((seq_length, 22), name=name, **kwargs)`
    """
    return Input((seq_length, len(AMINO_ACIDS)), name=name, **kwargs)

def InputRNAStructure(seq_length, name=None, **kwargs):
    """Input placeholder for array returned by `encodeRNAStructure`

    Wrapper for: `keras.layers.Input((seq_length, 5), name=name, **kwargs)`
    """
    return Input((seq_length, len(RNAplfold_PROFILES)), name=name, **kwargs)

# deprecated
def InputSplines(seq_length, n_bases=10, name=None, **kwargs):
    """Input placeholder for array returned by `encodeSplines`

    Wrapper for: `keras.layers.Input((seq_length, n_bases), name=name, **kwargs)`
    """
    return Input((seq_length, n_bases), name=name, **kwargs)

def InputSplines1D(seq_length, n_bases=10, name=None, **kwargs):
    """Input placeholder for array returned by `encodeSplines`

    Wrapper for: `keras.layers.Input((seq_length, n_bases), name=name, **kwargs)`
    """
    return Input((seq_length, n_bases), name=name, **kwargs)

# TODO - deprecate
def InputDNAQuantity(seq_length, n_features=1, name=None, **kwargs):
    """Convenience wrapper around `keras.layers.Input`:

    `Input((seq_length, n_features), name=name, **kwargs)`
    """
    return Input((seq_length, n_features), name=name, **kwargs)

# TODO - deprecate
def InputDNAQuantitySplines(seq_length, n_bases=10, name="DNASmoothPosition", **kwargs):
    """Convenience wrapper around keras.layers.Input:

    `Input((seq_length, n_bases), name=name, **kwargs)`
    """
    return Input((seq_length, n_bases), name=name, **kwargs)

# --------------------------------------------

class GlobalSumPooling1D(_GlobalPooling1D):
    """Global average pooling operation for temporal data.

    # Note
      - Input shape: 3D tensor with shape: `(batch_size, steps, features)`.
      - Output shape: 2D tensor with shape: `(batch_size, channels)`

    """

    def call(self, inputs):
        return K.sum(inputs, axis=1)

class ConvSequence(Conv1D):
    """Convenience wrapper over `keras.layers.Conv1D` with 3 changes:

    - additional plotting method: `plot_weights(index=None, plot_type="motif_raw", figsize=None, ncol=1)`
            - **index**: can be a particular index or a list of indicies
            - **plot_type**: Can be one of `"heatmap"`, `"motif_raw"`, `"motif_pwm"` or `"motif_pwm_info"`.
            - **figsize**: tuple, Figure size
            - **ncol**: Number of axis columns
    - additional argument `seq_length` instead of `input_shape`
    - restriction in build method: `input_shape[-1]` needs to the match the vocabulary size

    Clasess `Conv*` all inherit from `ConvSequence` and define the corresponding vocabulary:

    - ConvDNA
    - ConvRNA
    - ConvRNAStructure
    - ConvAA
    - ConvCodon
    """

    VOCAB = DNA

    def __init__(self,
                 filters,
                 kernel_size,
                 strides=1,
                 padding='valid',
                 dilation_rate=1,
                 activation=None,
                 use_bias=True,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 seq_length=None,
                 **kwargs):

        # override input shape
        if seq_length:
            kwargs["input_shape"] = (seq_length, len(self.VOCAB))
            kwargs.pop("batch_input_shape", None)

        super(ConvSequence, self).__init__(
            filters=filters,
            kernel_size=kernel_size,
            strides=strides,
            padding=padding,
            dilation_rate=dilation_rate,
            activation=activation,
            use_bias=use_bias,
            kernel_initializer=kernel_initializer,
            bias_initializer=bias_initializer,
            kernel_regularizer=kernel_regularizer,
            bias_regularizer=bias_regularizer,
            activity_regularizer=activity_regularizer,
            kernel_constraint=kernel_constraint,
            bias_constraint=bias_constraint,
            **kwargs)

        self.seq_length = seq_length

    def build(self, input_shape):
        if input_shape[-1] is not len(self.VOCAB):
            raise ValueError("{cls} requires input_shape[-1] == {n}. Given: {s}".
                             format(cls=self.__class__.__name__, n=len(self.VOCAB), s=input_shape[-1]))
        return super(ConvSequence, self).build(input_shape)

    def get_config(self):
        config = super(ConvSequence, self).get_config()
        config["seq_length"] = self.seq_length
        return config

    def _plot_weights_heatmap(self, index=None, figsize=None, **kwargs):
        """Plot weights as a heatmap

        index = can be a particular index or a list of indicies
        **kwargs - additional arguments to concise.utils.plot.heatmap
        """
        W = self.get_weights()[0]
        if index is None:
            index = np.arange(W.shape[2])

        fig = heatmap(np.swapaxes(W[:, :, index], 0, 1), plot_name="filter: ",
                      vocab=self.VOCAB, figsize=figsize, **kwargs)
        # plt.show()
        return fig

    def _plot_weights_motif(self, index, plot_type="motif_raw",
                            background_probs=DEFAULT_BASE_BACKGROUND,
                            ncol=1,
                            figsize=None):
        """Index can only be a single int
        """

        w_all = self.get_weights()

        if len(w_all) == 0:
            raise Exception("Layer needs to be initialized first")
        W = w_all[0]
        if index is None:
            index = np.arange(W.shape[2])

        if isinstance(index, int):
            index = [index]
        fig = plt.figure(figsize=figsize)

        if plot_type == "motif_pwm" and plot_type in self.AVAILABLE_PLOTS:
            arr = pssm_array2pwm_array(W, background_probs)
        elif plot_type == "motif_raw" and plot_type in self.AVAILABLE_PLOTS:
            arr = W
        elif plot_type == "motif_pwm_info" and plot_type in self.AVAILABLE_PLOTS:
            quasi_pwm = pssm_array2pwm_array(W, background_probs)
            arr = _pwm2pwm_info(quasi_pwm)
        else:
            raise ValueError("plot_type needs to be from {0}".format(self.AVAILABLE_PLOTS))

        fig = seqlogo_fig(arr, vocab=self.VOCAB_name, figsize=figsize, ncol=ncol, plot_name="filter: ")
        # fig.show()
        return fig

    def plot_weights(self, index=None, plot_type="motif_raw", figsize=None, ncol=1, **kwargs):
        """Plot filters as heatmap or motifs

        index = can be a particular index or a list of indicies
        **kwargs - additional arguments to concise.utils.plot.heatmap
        """

        if "heatmap" in self.AVAILABLE_PLOTS and plot_type == "heatmap":
            return self._plot_weights_heatmap(index=index, figsize=figsize, ncol=ncol, **kwargs)
        elif plot_type[:5] == "motif":
            return self._plot_weights_motif(index=index, plot_type=plot_type, figsize=figsize, ncol=ncol, **kwargs)
        else:
            raise ValueError("plot_type needs to be from {0}".format(self.AVAILABLE_PLOTS))

class ConvDNA(ConvSequence):
    VOCAB = DNA
    VOCAB_name = "DNA"
    AVAILABLE_PLOTS = ["heatmap", "motif_raw", "motif_pwm", "motif_pwm_info"]

class ConvRNA(ConvDNA):
    VOCAB = RNA
    VOCAB_name = "RNA"
    AVAILABLE_PLOTS = ["heatmap", "motif_raw", "motif_pwm", "motif_pwm_info"]

class ConvAA(ConvDNA):
    VOCAB = AMINO_ACIDS
    VOCAB_name = "AA"
    AVAILABLE_PLOTS = ["heatmap", "motif_raw"]

class ConvRNAStructure(ConvDNA):
    VOCAB = RNAplfold_PROFILES
    VOCAB_name = "RNAStruct"
    AVAILABLE_PLOTS = ["heatmap", "motif_raw"]

class ConvCodon(ConvSequence):
    VOCAB = CODONS
    AVAILABLE_PLOTS = ["heatmap"]

    def build(self, input_shape):
        if input_shape[-1] not in [len(CODONS), len(CODONS + STOP_CODONS)]:
            raise ValueError("{cls} requires input_shape[-1] == {n} or {m}".
                             format(cls=self.__class__.__name__,
                                    n=len(CODONS),
                                    m=len(CODONS + STOP_CODONS)))

        if input_shape[-1] == len(CODONS + STOP_CODONS):
            self.VOCAB = CODONS + STOP_CODONS

        return super(ConvSequence, self).build(input_shape)

# --------------------------------------------

############################################
# Smoothing layers

# TODO - re-write SplineWeight1D with SplineT layer
# TODO - SplineWeight1D - use new API and update
#        - think how to call share_splines...?
#        - use a regularizer rather than just
class SplineWeight1D(Layer):
    """Up- or down-weight positions in the activation array of 1D convolutions:

    `x^{out}_{ijk} = x^{in}_{ijk}* (1 + f_S^k(j)) \;,`
    where f_S is the spline transformation.

    # Arguments
        n_bases: int; Number of spline bases used for the positional effect.
        l2_smooth: (float) L2 regularization strength for the second
    order differences in positional bias' smooth splines. (GAM smoothing regularization)
        l2: (float) L2 regularization strength for the spline base coefficients.
        use_bias: boolean; should we add a bias to the transition
        bias_initializer: bias initializer - from `keras.initializers`
    """

    def __name__(self):
        return "SplineWeight1D"

    def __init__(self,
                 # spline type
                 n_bases=10,
                 spline_degree=3,
                 share_splines=False,
                 # regularization
                 l2_smooth=0,
                 l2=0,
                 use_bias=False,
                 bias_initializer='zeros',
                 **kwargs):
        self.n_bases = n_bases
        self.spline_degree = spline_degree
        self.share_splines = share_splines
        self.l2 = l2
        self.l2_smooth = l2_smooth
        self.use_bias = use_bias
        self.bias_initializer = initializers.get(bias_initializer)

        super(SplineWeight1D, self).__init__(**kwargs)

    def build(self, input_shape):
        # input_shape = (None, steps, filters)

        start = 0
        end = input_shape[1]
        filters = input_shape[2]

        if self.share_splines:
            n_spline_tracks = 1
        else:
            n_spline_tracks = filters

        # setup the bspline object
        self.bs = BSpline(start, end - 1,
                          n_bases=self.n_bases,
                          spline_order=self.spline_degree
                          )

        # create X_spline,
        self.positions = np.arange(end)
        self.X_spline = self.bs.predict(self.positions, add_intercept=False)  # shape = (end, self.n_bases)

        # convert to the right precision and K.constant
        self.X_spline_K = K.constant(K.cast_to_floatx(self.X_spline))

        # add weights - all set to 0
        self.kernel = self.add_weight(shape=(self.n_bases, n_spline_tracks),
                                      initializer='zeros',
                                      name='kernel',
                                      regularizer=GAMRegularizer(self.n_bases, self.spline_degree,
                                                                 self.l2_smooth, self.l2),
                                      trainable=True)

        if self.use_bias:
            self.bias = self.add_weight((n_spline_tracks, ),
                                        initializer=self.bias_initializer,
                                        name='bias',
                                        regularizer=None)

        super(SplineWeight1D, self).build(input_shape)  # Be sure to call this somewhere!

    def call(self, x):

        spline_track = K.dot(self.X_spline_K, self.kernel)

        if self.use_bias:
            spline_track = K.bias_add(spline_track, self.bias)

        # if self.spline_exp:
        #     spline_track = K.exp(spline_track)
        # else:
        spline_track = spline_track + 1

        # multiply together the two coefficients
        output = spline_track * x

        return output

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = {
            'n_bases': self.n_bases,
            'spline_degree': self.spline_degree,
            'share_splines': self.share_splines,
            # 'spline_exp': self.spline_exp,
            'l2_smooth': self.l2_smooth,
            'l2': self.l2,
            'use_bias': self.use_bias,
            'bias_initializer': initializers.serialize(self.bias_initializer),
        }
        base_config = super(SplineWeight1D, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def positional_effect(self):
        w = self.get_weights()[0]
        pos_effect = np.dot(self.X_spline, w)
        return {"positional_effect": pos_effect, "positions": self.positions}

    def plot(self, *args, **kwargs):
        pe = self.positional_effect()
        plt.plot(pe["positions"], pe["positional_effect"], *args, **kwargs)
        plt.xlabel("Position")
        plt.ylabel("Positional effect")

# SplineT -> use just locally-connected layers
# SplineT1D -> wrap ConvSplines -> However,
# It's better if we use just matrix multiplications.
# - We can have more control over the inputs
# - I think it's better to explicitly state the dimensions: 1D or so.

class SplineT(Layer):
    """Spline transformation layer.

    As input, it needs an array of scalars pre-processed by `concise.preprocessing.EncodeSplines`
    Specifically, the input/output dimensions are:

    - Input: N x ... x channels x n_bases
    - Output: N x ... x channels

    # Arguments
        shared_weights: bool, if True spline transformation weights
    are shared across different features
        kernel_regularizer: use `concise.regularizers.SplineSmoother`
        other arguments: See `keras.layers.Dense`
    """

    def __init__(self,
                 # regularization
                 shared_weights=False,
                 kernel_regularizer=None,
                 use_bias=False,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 **kwargs
                 ):
        super(SplineT, self).__init__(**kwargs)  # Be sure to call this somewhere!

        self.shared_weights = shared_weights
        self.kernel_regularizer = regularizers.get(kernel_regularizer)
        self.use_bias = use_bias
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)

        self.input_spec = InputSpec(min_ndim=3)

    def build(self, input_shape):
        assert len(input_shape) >= 3

        n_bases = input_shape[-1]
        n_features = input_shape[-2]

        # self.input_shape = input_shape
        self.inp_shape = input_shape
        self.n_features = n_features
        self.n_bases = n_bases

        if self.shared_weights:
            use_n_features = 1
        else:
            use_n_features = self.n_features

        # print("n_bases: {0}".format(n_bases))
        # print("n_features: {0}".format(n_features))

        self.kernel = self.add_weight(shape=(n_bases, use_n_features),
                                      initializer=self.kernel_initializer,
                                      name='kernel',
                                      regularizer=self.kernel_regularizer,
                                      trainable=True)

        if self.use_bias:
            self.bias = self.add_weight((n_features, ),
                                        initializer=self.bias_initializer,
                                        name='bias',
                                        regularizer=None)

        self.built = True
        super(SplineT, self).build(input_shape)  # Be sure to call this somewhere!

    def compute_output_shape(self, input_shape):
        return input_shape[:-1]

    def call(self, inputs):
        N = len(self.inp_shape)
        # put -2 axis (features) to the front
        # import pdb
        # pdb.set_trace()

        if self.shared_weights:
            return K.squeeze(K.dot(inputs, self.kernel), -1)

        output = K.permute_dimensions(inputs, (N - 2, ) + tuple(range(N - 2)) + (N - 1,))

        output_reshaped = K.reshape(output, (self.n_features, -1, self.n_bases))
        bd_output = K.batch_dot(output_reshaped, K.transpose(self.kernel))
        output = K.reshape(bd_output, (self.n_features, -1) + self.inp_shape[1:(N - 2)])
        # move axis 0 (features) to back
        output = K.permute_dimensions(output, tuple(range(1, N - 1)) + (0,))
        if self.use_bias:
            output = K.bias_add(output, self.bias, data_format="channels_last")
        return output

    def get_config(self):
        config = {
            'shared_weights': self.shared_weights,
            'kernel_regularizer': regularizers.serialize(self.kernel_regularizer),
            'use_bias': self.use_bias,
            'kernel_initializer': initializers.serialize(self.kernel_initializer),
            'bias_initializer': initializers.serialize(self.bias_initializer)
        }
        base_config = super(SplineT, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    # TODO - add X_spline as non-trainable weights

# Deprecated
class GAMSmooth(Layer):

    def __name__(self):
        return "GAMSmooth"

    def __init__(self,
                 # spline type
                 n_bases=10,
                 spline_order=3,
                 share_splines=False,
                 spline_exp=False,
                 # regularization
                 l2_smooth=1e-5,
                 l2=1e-5,
                 use_bias=False,
                 bias_initializer='zeros',
                 **kwargs):
        """

        # Arguments
            n_splines int: Number of splines used for the positional bias.
            spline_exp (bool): If True, the positional bias score is observed by: `np.exp(spline_score)`,
               where `spline_score` is the linear combination of B-spline basis functions.
               If False, `np.exp(spline_score + 1)` is used.
            l2 (float): L2 regularization strength for the second order differences in positional bias' smooth splines.
            (GAM smoothing regularization)
            l2_smooth (float): L2 regularization strength for the spline base coefficients.
            use_bias: boolean; should we add a bias to the transition
            bias_initializer; bias initializer - from keras.initailizers
        """
        self.n_bases = n_bases
        self.spline_order = spline_order
        self.share_splines = share_splines
        self.spline_exp = spline_exp
        self.l2 = l2
        self.l2_smooth = l2_smooth
        self.use_bias = use_bias
        self.bias_initializer = initializers.get(bias_initializer)

        super(GAMSmooth, self).__init__(**kwargs)

    def build(self, input_shape):
        # input_shape = (None, steps, filters)

        start = 0
        end = input_shape[1]
        filters = input_shape[2]

        if self.share_splines:
            n_spline_tracks = 1
        else:
            n_spline_tracks = filters

        # setup the bspline object
        self.bs = BSpline(start, end - 1,
                          n_bases=self.n_bases,
                          spline_order=self.spline_order
                          )

        # create X_spline,
        self.positions = np.arange(end)
        self.X_spline = self.bs.predict(self.positions, add_intercept=False)  # shape = (end, self.n_bases)

        # convert to the right precision and K.constant
        self.X_spline_K = K.constant(K.cast_to_floatx(self.X_spline))

        # add weights - all set to 0
        self.kernel = self.add_weight(shape=(self.n_bases, n_spline_tracks),
                                      initializer='zeros',
                                      name='kernel',
                                      regularizer=GAMRegularizer(self.n_bases, self.spline_order,
                                                                 self.l2_smooth, self.l2),
                                      trainable=True)

        if self.use_bias:
            self.bias = self.add_weight((n_spline_tracks, ),
                                        initializer=self.bias_initializer,
                                        name='bias',
                                        regularizer=None)

        super(GAMSmooth, self).build(input_shape)  # Be sure to call this somewhere!

    def call(self, x):

        spline_track = K.dot(self.X_spline_K, self.kernel)

        if self.use_bias:
            spline_track = K.bias_add(spline_track, self.bias)

        if self.spline_exp:
            spline_track = K.exp(spline_track)
        else:
            spline_track = spline_track + 1

        # multiply together the two coefficients
        output = spline_track * x

        return output

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = {
            'n_bases': self.n_bases,
            'spline_order': self.spline_order,
            'share_splines': self.share_splines,
            'spline_exp': self.spline_exp,
            'l2_smooth': self.l2_smooth,
            'l2': self.l2,
            'use_bias': self.use_bias,
            'bias_initializer': initializers.serialize(self.bias_initializer),
        }
        base_config = super(GAMSmooth, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def positional_effect(self):
        w = self.get_weights()[0]
        pos_effect = np.dot(self.X_spline, w)
        return {"positional_effect": pos_effect, "positions": self.positions}

    def plot(self, *args, **kwargs):
        pe = self.positional_effect()
        plt.plot(pe["positions"], pe["positional_effect"], *args, **kwargs)
        plt.xlabel("Position")
        plt.ylabel("Positional effect")

# ResSplineWeight
# SplineWeight ?
# WeightedSum1D

# TODO - add the plotting functionality
# TODO - rename the layer
# additional arguments?
# - share_splines=False,
# - spline_exp=False
#
# TODO - use similar arguments to GAMSmooth (not as a thin wrapper around Conv1d)
# TODO - fix & unit-test this layer

# ConvSplineTr1D
# DenseSplineTr
# SplineTr
class ConvSplines(Conv1D):
    """Convenience wrapper over `keras.layers.Conv1D` with 2 changes:
    - additional argument seq_length specifying input_shape (as in ConvDNA)
    - restriction in kernel_regularizer - needs to be of class GAMRegularizer
    - hard-coded values:
       - kernel_size=1,
       - strides=1,
       - padding='valid',
       - dilation_rate=1,
    """

    def __init__(self, filters,
                 activation=None,
                 use_bias=True,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer=GAMRegularizer(),
                 bias_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 activity_regularizer=None,
                 **kwargs):

        super(ConvSplines, self).__init__(
            filters=filters,
            kernel_size=1,
            strides=1,
            padding='valid',
            dilation_rate=1,
            activation=activation,
            use_bias=use_bias,
            kernel_initializer=kernel_initializer,
            bias_initializer=bias_initializer,
            kernel_regularizer=kernel_regularizer,
            bias_regularizer=bias_regularizer,
            activity_regularizer=activity_regularizer,
            kernel_constraint=kernel_constraint,
            bias_constraint=bias_constraint,
            **kwargs)

        if not isinstance(self.kernel_regularizer, regularizers.GAMRegularizer):
            raise ValueError("Regularizer has to be of type concise.regularizers.GAMRegularizer. " +
                             "Current type: " + str(type(self.kernel_regularizer)),
                             "\nObject: " + str(self.kernel_regularizer))

        # self.seq_length = seq_length

    def build(self, input_shape):
        # update the regularizer
        self.kernel_regularizer.n_bases = input_shape[2]

        return super(ConvSplines, self).build(input_shape)

    def get_config(self):
        config = super(ConvSplines, self).get_config()
        config.pop('kernel_size')
        config.pop('strides')
        config.pop('padding')
        config.pop('dilation_rate')
        # config["seq_length"] = self.seq_length
        return config

class BiDropout(Dropout):
    """Applies Dropout to the input, no matter if in learning phase or not.
    """

    def __init__(self, bi_dropout=True, **kwargs):
        # __init__(self, rate, noise_shape=None, seed=None, **kwargs)
        super(BiDropout, self).__init__(**kwargs)
        self.bi_dropout = bi_dropout

    #
    def call(self, inputs, training=None):
        if 0. < self.rate < 1.:
            noise_shape = self._get_noise_shape(inputs)

            #
            def dropped_inputs():
                return K.dropout(inputs, self.rate, noise_shape, seed=self.seed)

            if self.bi_dropout:
                # K.in_train_phase returns the first argument if in training phase otherwise the second
                # return K.in_train_phase(dropped_inputs, inputs, training=training)
                # Taken from keras.backend.tensorflow_backend
                if callable(dropped_inputs):
                    return dropped_inputs()
                else:
                    return dropped_inputs
            else:
                return K.in_train_phase(dropped_inputs, inputs,
                                        training=training)
        return inputs

    #
    @classmethod
    def create_from_dropout(cls, dropout_obj):
        if not isinstance(dropout_obj, Dropout):
            raise Exception("Only Dropout objects can be converted this way!")
        kwargs = dropout_obj.get_config()
        # alternatively can we use "get_config" in combination with (Layer.__init__)allowed_kwargs?
        return cls(**kwargs)

# backcompatibility
ConvDNAQuantitySplines = ConvSplines

AVAILABLE = ["InputDNA", "ConvDNA",
             "InputRNA", "ConvRNA",
             "InputCodon", "ConvCodon",
             "InputAA", "ConvAA",
             "InputRNAStructure", "ConvRNAStructure",
             "InputSplines", "ConvSplines",
             "GlobalSumPooling1D",
             "SplineWeight1D",
             "SplineT",
             # legacy
             "InputDNAQuantitySplines", "InputDNAQuantity",
             "GAMSmooth", "ConvDNAQuantitySplines", "BiDropout"]

def get(name):
    return get_from_module(name, globals())


import torch
from torch.nn import functional as F

from mmdet3d.ops import Voxelization
from mmdet.models import DETECTORS
from .. import builder
from .two_stage import TwoStage3DDetector
from .voxelnet import Fusion
import numpy as np

@DETECTORS.register_module()
class PartA2(TwoStage3DDetector):
    r"""Part-A2 detector.

    Please refer to the `paper <https://arxiv.org/abs/1907.03670>`_
    """

    def __init__(self,
                 voxel_layer,
                 voxel_encoder,
                 middle_encoder,
                 backbone,
                 neck=None,
                 rpn_head=None,
                 roi_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None):
        super(PartA2, self).__init__(
            backbone=backbone,
            neck=neck,
            rpn_head=rpn_head,
            roi_head=roi_head,
            train_cfg=train_cfg,
            test_cfg=test_cfg,
            pretrained=pretrained,
        )
        self.voxel_layer = Voxelization(**voxel_layer)
        self.voxel_encoder = builder.build_voxel_encoder(voxel_encoder)
        self.middle_encoder = builder.build_middle_encoder(middle_encoder)

        # 转换RangeImage部分相关代码
        self.H = 48
        self.W = 512
        self.fov_up = 3
        self.fov_down = -15.0
        self.pi = torch.tensor(np.pi)
        fov_up = self.fov_up * self.pi / 180.0
        fov_down = self.fov_down * self.pi / 180.0
        fov = abs(fov_up) + abs(fov_down)
        self.uv = torch.zeros((2, self.H, self.W))
        self.uv[1] = torch.arange(0, self.W)
        self.uv.permute((0, 2, 1))[0] = torch.arange(0, self.H)
        self.uv[0] = ((self.H - self.uv[0]) * fov - abs(fov_down) * self.H) / self.H
        self.uv[1] = (self.uv[1] * 2.0 - self.W) * self.pi / (self.W * 4)  # 最后一个 4 用来控制水平范围

        # self.range_encoder = RangeEncoder(5, 64, use_img=True)
        self.fusion = Fusion(5, 3)

    def extract_feat(self, points, img_metas):
        """Extract features from points."""
        voxel_dict = self.voxelize(points)
        voxel_features = self.voxel_encoder(voxel_dict['voxels'],
                                            voxel_dict['num_points'],
                                            voxel_dict['coors'])
        batch_size = voxel_dict['coors'][-1, 0].item() + 1
        feats_dict = self.middle_encoder(voxel_features, voxel_dict['coors'],
                                         batch_size)
        x = self.backbone(feats_dict['spatial_features'])
        if self.with_neck:
            neck_feats = self.neck(x)
            feats_dict.update({'neck_feats': neck_feats})
        return feats_dict, voxel_dict

    @torch.no_grad()
    def voxelize(self, points):
        """Apply hard voxelization to points."""
        voxels, coors, num_points, voxel_centers = [], [], [], []
        for res in points:
            res_voxels, res_coors, res_num_points = self.voxel_layer(res)
            res_voxel_centers = (
                res_coors[:, [2, 1, 0]] + 0.5) * res_voxels.new_tensor(
                    self.voxel_layer.voxel_size) + res_voxels.new_tensor(
                        self.voxel_layer.point_cloud_range[0:3])
            voxels.append(res_voxels)
            coors.append(res_coors)
            num_points.append(res_num_points)
            voxel_centers.append(res_voxel_centers)

        voxels = torch.cat(voxels, dim=0)
        num_points = torch.cat(num_points, dim=0)
        voxel_centers = torch.cat(voxel_centers, dim=0)
        coors_batch = []
        for i, coor in enumerate(coors):
            coor_pad = F.pad(coor, (1, 0), mode='constant', value=i)
            coors_batch.append(coor_pad)
        coors_batch = torch.cat(coors_batch, dim=0)

        voxel_dict = dict(
            voxels=voxels,
            num_points=num_points,
            coors=coors_batch,
            voxel_centers=voxel_centers)
        return voxel_dict

    def forward_train(self,
                      points,
                      img_metas,
                      gt_bboxes_3d,
                      gt_labels_3d,
                      img=None,
                      gt_bboxes_ignore=None,
                      proposals=None):
        """Training forward function.

        Args:
            points (list[torch.Tensor]): Point cloud of each sample.
            img_metas (list[dict]): Meta information of each sample
            gt_bboxes_3d (list[:obj:`BaseInstance3DBoxes`]): Ground truth
                boxes for each sample.
            gt_labels_3d (list[torch.Tensor]): Ground truth labels for
                boxes of each sampole
            gt_bboxes_ignore (list[torch.Tensor], optional): Ground truth
                boxes to be ignored. Defaults to None.

        Returns:
            dict: Losses of each branch.
        """

        # 转换成range提取特征后转回lidar
        batchsize = len(points)
        rangeImage = []
        for i in range(batchsize):
            rangeImage.append(self.lidar_to_range_gpu(points[i]).unsqueeze(0))
        rangeImage = torch.cat(rangeImage, dim=0)
        # 是否加入img信息
        # range_feat = self.range_encoder(rangeImage, img)      #用自编码器的形式
        range_feat = self.fusion(rangeImage, img)
        range_ori = torch.cat((rangeImage[:, 0:2], range_feat), dim=1)
        pts_with_range = []
        for i in range(batchsize):
            pts_with_range.append(self.range_to_lidar_gpu(range_ori[i].squeeze(0)))

        feats_dict, voxels_dict = self.extract_feat(pts_with_range, img_metas)      #point 换成pts_with_range

        losses = dict()

        if self.with_rpn:
            rpn_outs = self.rpn_head(feats_dict['neck_feats'])
            rpn_loss_inputs = rpn_outs + (gt_bboxes_3d, gt_labels_3d,
                                          img_metas)
            rpn_losses = self.rpn_head.loss(
                *rpn_loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
            losses.update(rpn_losses)

            proposal_cfg = self.train_cfg.get('rpn_proposal',
                                              self.test_cfg.rpn)
            proposal_inputs = rpn_outs + (img_metas, proposal_cfg)
            proposal_list = self.rpn_head.get_bboxes(*proposal_inputs)
        else:
            proposal_list = proposals

        roi_losses = self.roi_head.forward_train(feats_dict, voxels_dict,
                                                 img_metas, proposal_list,
                                                 gt_bboxes_3d, gt_labels_3d)

        losses.update(roi_losses)

        return losses

    def simple_test(self, points, img_metas, imgs=None, proposals=None, rescale=False):
        """Test function without augmentaiton."""

        # 转换成range提取特征后转回lidar
        batchsize = len(points)
        rangeImage = []
        for i in range(batchsize):
            rangeImage.append(self.lidar_to_range_gpu(points[i]).unsqueeze(0))
        rangeImage = torch.cat(rangeImage, dim=0)
        # 是否加入img信息
        # range_feat = self.range_encoder(rangeImage, img)      #用自编码器的形式
        range_feat = self.fusion(rangeImage, imgs)
        range_ori = torch.cat((rangeImage[:, 0:2], range_feat), dim=1)
        pts_with_range = []
        for i in range(batchsize):
            pts_with_range.append(self.range_to_lidar_gpu(range_ori[i].squeeze(0)))

        feats_dict, voxels_dict = self.extract_feat(pts_with_range, img_metas)

        if self.with_rpn:
            rpn_outs = self.rpn_head(feats_dict['neck_feats'])
            proposal_cfg = self.test_cfg.rpn
            bbox_inputs = rpn_outs + (img_metas, proposal_cfg)
            proposal_list = self.rpn_head.get_bboxes(*bbox_inputs)
        else:
            proposal_list = proposals

        return self.roi_head.simple_test(feats_dict, voxels_dict, img_metas,
                                         proposal_list)

    def lidar_to_range_gpu(self, points):
        device = points.device
        pi = torch.tensor(np.pi).to(device)
        fov_up = self.fov_up * pi / 180.0
        fov_down = self.fov_down * pi / 180.0
        fov = abs(fov_up) + abs(fov_down)

        depth = torch.norm(points, 2, dim=1)

        x = points[:, 0]
        y = points[:, 1]
        z = points[:, 2]

        yaw = torch.atan2(y, x)
        pitch = torch.asin(z / depth)

        u = 0.5 * (1 - 4 * yaw / pi) * self.W  # 最后一个 4 用来控制水平范围
        v = (1 - (pitch + abs(fov_down)) / fov) * self.H

        zero_tensor = torch.zeros_like(u)
        W_tensor = torch.ones_like(u) * (self.W - 1)
        H_tensor = torch.ones_like(v) * (self.H - 1)

        u = torch.floor(u)
        u = torch.min(u, W_tensor)
        u = torch.max(u, zero_tensor).long()

        v = torch.floor(v)
        v = torch.min(v, H_tensor)
        v = torch.max(v, zero_tensor).long()

        range_image = torch.full((5, self.H, self.W), 0, dtype=torch.float32).to(device)
        range_image[0][v, u] = depth
        range_image[1][v, u] = points[:, 3]
        range_image[2][v, u] = points[:, 0]
        range_image[3][v, u] = points[:, 1]
        range_image[4][v, u] = points[:, 2]
        return range_image

    def range_to_lidar_gpu(self, range_img):
        device = range_img.device
        self.uv = self.uv.to(device)
        lidar_out = torch.zeros((12, self.H, self.W)).to(device)
        lidar_out[0] = range_img[0] * torch.cos(self.uv[0]) * torch.cos(self.uv[1])
        lidar_out[1] = range_img[0] * torch.cos(self.uv[0]) * torch.sin(self.uv[1]) * (-1)
        lidar_out[2] = range_img[0] * torch.sin(self.uv[0])
        lidar_out[3:] = range_img[1:]
        lidar_out = lidar_out.permute((2, 1, 0)).reshape([-1, 12])
        lidar_out = lidar_out[torch.where(lidar_out[:, 0] != 0)]
        return lidar_out


from PIL import Image, ImageDraw, ImageFilter
from random import randint
from blurEdges import blurImage
from yolo3.utils import get_random_data
import csv
import numpy as np
import os
import glob

from matplotlib import pyplot as plt

def cutNGetImg(annotation_line,prefix):
    '''random preprocessing for real-time data augmentation'''
    line = annotation_line.split()
    image = Image.open(prefix + line[0])
    box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])
    image = np.array(image)/np.max(np.array(image))
    image_data = list()
    class_list = box[:, 4]
    box=box[:, :4]
    if (len(image.shape)==2):
        for b in range(len(box)):
            box1 = box[b, :]
            image_data.append(image[box1[1]:box1[3], box1[0]:box1[2]])
    else:
        for b in range(len(box)):
            box1 = box[b,:]
            image_data.append(image[box1[1]:box1[3], box1[0]:box1[2], :])

    return image_data, class_list

#TODO: Choose nr of images, blur images, define cooridnates and define what to write to annotation if num_sign=0

prefix = 'C:/Users/CTK_CAD/Chalmers Teknologkonsulter AB/Bird Classification - Images/Bird Detection - Images/Original images/'

road_path = 'C:\\Users\\CTK_CAD\\PycharmProjects\\BirdDetection\\keras-yolo3\\test_background\\' # background
#road_folder = os.listdir(road_path)  # create a folder to iterate through

aug_path = prefix + 'test/' # Final destination

'''
files = glob.glob(aug_path+'*.png')
for f in files:
    os.remove(f)

'''

road_size = (1200,900)
start = 501
max = start + 100 #number of images
number_of_roads = 20
num_classes = 3

# plt.bar(range(len(dist)), dist, 1/1.5, color="blue")
# plt.show()      #Plot distribution

for i in range(start,max):
    text_file = open(aug_path + "test_annotations.txt", "a")
    road_num = randint(1,number_of_roads) #Random road 'road_num.png'
    road = Image.open(road_path+str(road_num)+'.jpg')
    road = road.resize(road_size)
    num_of_signs_in_img = randint(1,3) #number of signs to be pasted into the image
    box_str = ''
    for j in range(0, num_of_signs_in_img):
        class_nr = np.random.choice(num_classes)  #chooses a class
        annotations_path = prefix + str(class_nr) + '/' + 'annotations' + str(class_nr) + '.txt'
        '''
        gtFile = open(annotations_path)  # annotations file
        gtReader = csv.reader(gtFile, delimiter=';')  # csv parser for annotations file
        next(gtReader)  # skip header
        '''
        #row_num = randint(1, sum(1 for row in gtReader))

        with open(annotations_path) as f:  # annotations file
            gtReader = f.readlines()  # read the txt annotations file

        a = np.random.choice(len(gtReader))

        sign, class_list = cutNGetImg(gtReader[a],prefix)

            #print(sign)
        for c, img in enumerate(sign):
            #print(img)

            if road_size[0]>img.shape[0] and road_size[1]>img.shape[1]:
                x = randint(0, road_size[0]-img.shape[1])
                y = randint(0, road_size[1]-img.shape[0])
                img = Image.fromarray(np.uint8(img*255))
                road = blurImage(road,img,x,y)
                box_str += ' ' + str(x)+','+str(y)+','+str(x+img.size[0])+','+str(y+img.size[1])+','+str(class_list[c])

    print('i:'+ str(i))
    text_file.write('augmented/'+format(i,'05d')+'.png'+box_str+'\n') #00600.png;774;411;815;446;11 in .txt-file, separate rows.

    road.save(aug_path + format(i, '05d')+'.png')
    text_file.close()


import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision.models as models
import numpy as np
from utils.config import cfg
from roi_pooling.modules.roi_pool import _RoIPooling
from roi_crop.modules.roi_crop import _RoICrop
from roi_align.modules.roi_align import RoIAlignAvg
from rpn.proposal_target_layer_cascade import _ProposalTargetLayer
import time
import pdb
from utils.net_utils import _smooth_l1_loss, _crop_pool_layer, _affine_grid_gen, _affine_theta
from rpn.bbox_transform import bbox_transform_inv, clip_boxes

class _fastRCNN(nn.Module):
    """ fast RCNN """

    def __init__(self, classes, class_agnostic):
        super(_fastRCNN, self).__init__()
        self.classes = classes
        self.n_classes = len(classes)
        self.class_agnostic = class_agnostic

        # loss
        self.RCNN_loss_cls = 0
        self.RCNN_loss_bbox = 0

        self.RCNN_proposal_target = _ProposalTargetLayer(self.n_classes)
        self.RCNN_roi_pool = _RoIPooling(
            cfg.POOLING_SIZE, cfg.POOLING_SIZE, 1.0 / 16.0)
        self.RCNN_roi_align = RoIAlignAvg(
            cfg.POOLING_SIZE, cfg.POOLING_SIZE, 1.0 / 16.0)

        self.grid_size = cfg.POOLING_SIZE * \
            2 if cfg.CROP_RESIZE_WITH_MAX_POOL else cfg.POOLING_SIZE
        self.RCNN_roi_crop = _RoICrop()

    def forward_rcnn_batch(self, base_feat, branch, rois, wgt_boxes, wnum_boxes, gt_boxes, num_boxes, im_info, image_classes, output_refine=False):
        batch_size = base_feat.size(0)

        # if it is training phrase, then use ground truth bboxes for refining
        if self.training:
            roi_data = self.RCNN_proposal_target(
                rois, wgt_boxes, wnum_boxes, gt_boxes, num_boxes)
            out_rois, rois_label, rois_target, rois_inside_ws, rois_outside_ws = roi_data

            rois_label = Variable(rois_label.view(-1).long())
            rois_target = Variable(rois_target.view(-1, rois_target.size(2)))
            rois_inside_ws = Variable(
                rois_inside_ws.view(-1, rois_inside_ws.size(2)))
            rois_outside_ws = Variable(
                rois_outside_ws.view(-1, rois_outside_ws.size(2)))
        else:
            out_rois = rois
            rois_label = None
            rois_target = None
            rois_inside_ws = None
            rois_outside_ws = None

        out_rois = Variable(out_rois)

        # do roi pooling based on predicted rois
        if cfg.POOLING_MODE == 'crop':
            # pdb.set_trace()
            # pooled_feat_anchor = _crop_pool_layer(base_feat, rois.view(-1, 5))
            grid_xy = _affine_grid_gen(
                out_rois.view(-1, 5), base_feat.size()[2:], self.grid_size)
            grid_yx = torch.stack(
                [grid_xy.data[:, :, :, 1], grid_xy.data[:, :, :, 0]], 3).contiguous()
            pooled_feat = self.RCNN_roi_crop(
                base_feat, Variable(grid_yx).detach())
            if cfg.CROP_RESIZE_WITH_MAX_POOL:
                pooled_feat = F.max_pool2d(pooled_feat, 2, 2)
        elif cfg.POOLING_MODE == 'align':
            pooled_feat = self.RCNN_roi_align(base_feat, out_rois.view(-1, 5))
        elif cfg.POOLING_MODE == 'pool':
            pooled_feat = self.RCNN_roi_pool(base_feat, out_rois.view(-1, 5))

        # feed pooled features to top model
        pooled_feat = self._head_to_tail(pooled_feat, branch)

        # compute bbox offset
        bbox_pred = branch.RCNN_bbox_pred(pooled_feat)
        if self.training and not self.class_agnostic:
            # select the corresponding columns according to roi labels
            bbox_pred_view = bbox_pred.view(
                bbox_pred.size(0), int(bbox_pred.size(1) / 4), 4)
            bbox_pred_select = torch.gather(bbox_pred_view, 1, rois_label.view(
                rois_label.size(0), 1, 1).expand(rois_label.size(0), 1, 4))
            bbox_pred = bbox_pred_select.squeeze(1)

        # compute object classification probability
        cls_score = branch.RCNN_cls_score(pooled_feat)
        cls_prob = F.softmax(cls_score, 1)

        RCNN_loss_cls = 0
        RCNN_loss_bbox = 0

        if self.training:

            # classification loss
            RCNN_loss_cls = F.cross_entropy(cls_score, rois_label)

            # bounding box regression L1 loss
            RCNN_loss_bbox = _smooth_l1_loss(
                bbox_pred, rois_target, rois_inside_ws, rois_outside_ws)

            # add image-level label regularization
            rois_batch_size = out_rois.size(1)
            rois_prob = F.softmax(cls_score, 1).view(batch_size, rois_batch_size, -1)

            valid_rois_prob = (rois_label > 0).view(batch_size, rois_batch_size, -1).float()
            rois_attention = F.softmax(cls_score, 1).view(batch_size, rois_batch_size, -1)
            rois_attention = rois_attention * valid_rois_prob

            # ignore background
            rois_prob = rois_prob[:, :, 1:]
            rois_attention = rois_attention[:, :, 1:]

            # rois_attention_prob = torch.sum(rois_prob * rois_attention, dim=1) / (torch.sum(rois_attention, dim=1) + 1e-10)
            rois_attention_prob, _ = torch.max(rois_prob, dim=1)
            image_loss_cls = F.binary_cross_entropy(rois_attention_prob, image_classes[:, 1:])
        else:
            image_loss_cls = None

        if self.training:
            cls_prob = cls_prob.view(batch_size, out_rois.size(1), -1)
            bbox_pred = bbox_pred.view(batch_size, out_rois.size(1), -1)
        else:
            cls_prob = cls_prob.view(1, out_rois.size(1), -1)
            bbox_pred = bbox_pred.view(1, out_rois.size(1), -1)

        if self.training and output_refine:
            # get transformation for wgt_boxes
            wgt_rois = wgt_boxes.new(wgt_boxes.size()).zero_()
            wgt_rois[:, :, 1:5] = wgt_boxes[:, :, :4]
            batch_size = base_feat.size(0)
            for i in range(batch_size):
                wgt_rois[:, :, 0] = i

            # do roi pooling based on predicted rois
            if cfg.POOLING_MODE == 'crop':
                # pdb.set_trace()
                # pooled_feat_anchor = _crop_pool_layer(base_feat, rois.view(-1, 5))
                grid_xy = _affine_grid_gen(
                    wgt_rois.view(-1, 5), base_feat.size()[2:], self.grid_size)
                grid_yx = torch.stack(
                    [grid_xy.data[:, :, :, 1], grid_xy.data[:, :, :, 0]], 3).contiguous()
                gt_pooled_feat = self.RCNN_roi_crop(
                    base_feat, Variable(grid_yx).detach())
                if cfg.CROP_RESIZE_WITH_MAX_POOL:
                    gt_pooled_feat = F.max_pool2d(gt_pooled_feat, 2, 2)
            elif cfg.POOLING_MODE == 'align':
                gt_pooled_feat = self.RCNN_roi_align(
                    base_feat, wgt_rois.view(-1, 5))
            elif cfg.POOLING_MODE == 'pool':
                gt_pooled_feat = self.RCNN_roi_pool(
                    base_feat, wgt_rois.view(-1, 5))

            # feed pooled features to top model
            gt_pooled_feat = self._head_to_tail(gt_pooled_feat, branch)

            # compute bbox offset
            wgt_bbox_delta = branch.RCNN_bbox_pred(gt_pooled_feat)
            wgt_bbox_delta = wgt_bbox_delta.view(-1, 4) * torch.FloatTensor(
                cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
            wgt_bbox_delta = wgt_bbox_delta.view(batch_size, -1, 4 * 21)
            wgt_bbox_out_rois = bbox_transform_inv(
                wgt_boxes, wgt_bbox_delta, batch_size)

            wgt_bbox_out_rois = clip_boxes(
                wgt_bbox_out_rois, im_info.data, batch_size)

            wgt_bbox_out = wgt_boxes.new(wgt_boxes.size()).zero_()

            wgt_cls = Variable(
                wgt_boxes[:, :, 4].data, requires_grad=False).long()
            for i in range(batch_size):
                for j in range(20):
                    cls_ind = wgt_cls[i, j]
                    wgt_bbox_out[i, j, :4] = wgt_bbox_out_rois[i,
                                                               j, cls_ind * 4:cls_ind * 4 + 4]

            wgt_bbox_out[:, :, 4] = wgt_boxes[:, :, 4]

            wgt_boxes_x = (wgt_boxes[:, :, 2] - wgt_boxes[:, :, 0] + 1)
            wgt_boxes_y = (wgt_boxes[:, :, 3] - wgt_boxes[:, :, 1] + 1)
            wgt_area_zero = (wgt_boxes_x == 1) & (wgt_boxes_y == 1)
            wgt_bbox_out.masked_fill_(wgt_area_zero.view(
                batch_size, wgt_area_zero.size(1), 1).expand(wgt_boxes.size()), 0)
            wgt_bbox_out = wgt_bbox_out.detach()
        else:
            wgt_bbox_out = None

        return (out_rois, cls_prob, bbox_pred, RCNN_loss_cls, RCNN_loss_bbox, rois_label, image_loss_cls), wgt_bbox_out

    def forward(self, im_data, im_info, gt_boxes, num_boxes, wgt_boxes, wnum_boxes, rois, image_classes):
        im_info = im_info.data
        gt_boxes = gt_boxes.data
        num_boxes = num_boxes.data
        wgt_boxes = wgt_boxes.data
        wnum_boxes = wnum_boxes.data
        rois = rois.data

        # feed image data to base model to obtain base feature map
        base_feat = self.RCNN_base(im_data)

        if self.training:
            # weak label branch and transformation regression
            rcnn_out_1, wgt_boxes_1 = self.forward_rcnn_batch(
                base_feat, self.wbranch, rois, wgt_boxes, wnum_boxes, gt_boxes, num_boxes,
                im_info, image_classes, output_refine=True)

            rcnn_out_2 = None

            # transformation label branch
            rcnn_out_3, wgt_boxes_2 = self.forward_rcnn_batch(
                base_feat, self.sbranch, rois, wgt_boxes_1, wnum_boxes, gt_boxes, num_boxes,
                im_info, image_classes, output_refine=True)

            # transformation label branch2
            rcnn_out_4, _ = self.forward_rcnn_batch(
                base_feat[:1], self.sbranch2, rois[:1], gt_boxes[:1], num_boxes[:1], gt_boxes[:1], num_boxes[:1],
                im_info[:1], image_classes[:1])

            rcnn_out_5, _ = self.forward_rcnn_batch(
                base_feat, self.sbranch2, rois, wgt_boxes_2, wnum_boxes, gt_boxes, num_boxes,
                im_info, image_classes)
        else:
            rcnn_out_1 = None
            rcnn_out_2 = None
            rcnn_out_3 = None
            rcnn_out_4 = None

            # transformation label branch
            rcnn_out_5, _ = self.forward_rcnn_batch(
                base_feat, self.sbranch2, rois, wgt_boxes, wnum_boxes, gt_boxes, num_boxes, im_info, image_classes)

        return rcnn_out_1, rcnn_out_2, rcnn_out_3, rcnn_out_4, rcnn_out_5

    def _init_weights(self, branch):
        def normal_init(m, mean, stddev, truncated=False):
            """
            weight initalizer: truncated normal and random normal.
            """
            # x is a parameter
            if truncated:
                m.weight.data.normal_().fmod_(2).mul_(stddev).add_(
                    mean)  # not a perfect approximation
            else:
                m.weight.data.normal_(mean, stddev)
                m.bias.data.zero_()

        normal_init(branch.RCNN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(branch.RCNN_bbox_pred, 0, 0.001, cfg.TRAIN.TRUNCATED)

    def create_architecture(self):
        self._init_modules()

        # init weak and strong branch
        self._init_weights(self.sbranch)
        self._init_weights(self.wbranch)
        self._init_weights(self.sbranch2)


from typing import Tuple

import glm
import imageio
from numpy.core.multiarray import ndarray

import pxng
from pxng import resource
from OpenGL.GL import GL_TRIANGLES, GL_RGBA, GL_RGB, glGenTextures, \
    GL_TEXTURE_RECTANGLE, glBindTexture, glTexParameteri, GL_TEXTURE_MAG_FILTER, \
    GL_NEAREST, GL_TEXTURE_MIN_FILTER, glTexImage2D, GL_UNSIGNED_BYTE, glTexSubImage2D, \
    GL_BLEND, glDisable, GL_RED, glPixelStorei, GL_UNPACK_ALIGNMENT, glGetInteger

class SpriteRectangle:
    def __init__(self):
        program = pxng.ShaderProgram('SpriteShader')
        program.add_shader(resource('shaders/sprite.vert'), pxng.ShaderType.Vertex)
        program.add_shader(resource('shaders/sprite.frag'), pxng.ShaderType.Fragment)
        program.compile_and_link()

        program.add_uniform('projection_view', glm.mat4x4)
        program.add_uniform('model', glm.mat4x4)
        program.add_uniform('texture_matrix', glm.mat4x4)
        program.add_uniform('color', glm.vec4)
        program.add_uniform('sprite_texture', glm.ivec1)
        self._program = program

        self._vao = pxng.VertexArrayObject(GL_TRIANGLES)
        self._vao.attach_buffer(pxng.BufferObject(data_type=glm.vec3))  # vertex buffer
        self._vao.attach_buffer(pxng.BufferObject(data_type=glm.uvec2))  # texture buffer

        self._vao.add_quad(
            glm.vec3(0, 0, 0),
            glm.vec3(0, -1, 0),
            glm.vec3(1, -1, 0),
            glm.vec3(1, 0, 0),
        )

        self._vao.set_texture(
            glm.uvec2(0, 1),
            glm.uvec2(0, 0),
            glm.uvec2(1, 0),
            glm.uvec2(1, 1)
        )

    def draw(self, spaces: pxng.Spaces):
        if self._vao.bind():
            self._program.activate()
            self._program.set_uniform('projection_view', spaces.projection_view)
            self._program.set_uniform('model', spaces.model.m)
            self._program.set_uniform('texture_matrix', spaces.texture.m)
            self._program.set_uniform('color', spaces.tint)
            self._program.set_uniform('sprite_texture', 0)
            self._vao.draw()

class Sprite:
    def __init__(self, data: ndarray):
        self._data = data
        self._created = False
        self._dirty = False
        self._texid = None
        self._width = data.shape[1]
        self._height = data.shape[0]
        self._components = 1
        if len(data.shape) > 2:
            self._components = data.shape[2]

        if self._components == 4:
            self._format = GL_RGBA
        elif self._components == 3:
            self._format = GL_RGB
        elif self._components == 1:
            self._format = GL_RED

        self._rect = None

    @property
    def width(self):
        return self._width

    @property
    def height(self):
        return self._height

    def set_pixel(self, x, y, color: Tuple[int, int, int]):
        self._data[y, x] = color
        self._dirty = True

    def _set_unpack_alignment(self):
        self._unpack_alignment = glGetInteger(GL_UNPACK_ALIGNMENT)
        if self._width % 2 == 1 or self._width % 4 == 2:
            # Odd sized texture
            glPixelStorei(GL_UNPACK_ALIGNMENT, 1)

    def _reset_unpack_alignment(self):
        glPixelStorei(GL_UNPACK_ALIGNMENT, self._unpack_alignment)

    def _create(self):
        self._texid = glGenTextures(1)
        glBindTexture(GL_TEXTURE_RECTANGLE, self._texid)

        glTexParameteri(GL_TEXTURE_RECTANGLE, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_RECTANGLE, GL_TEXTURE_MIN_FILTER, GL_NEAREST)

        d = self._data
        w = self._width
        h = self._height
        fmt = self._format

        self._set_unpack_alignment()
        glTexImage2D(GL_TEXTURE_RECTANGLE, 0, fmt, w, h, 0, fmt, GL_UNSIGNED_BYTE, d)
        self._reset_unpack_alignment()

        self._rect = SpriteRectangle()
        self._created = True

    def _update(self):
        glBindTexture(GL_TEXTURE_RECTANGLE, self._texid)

        d = self._data
        w = self._width
        h = self._height
        fmt = self._format

        self._set_unpack_alignment()
        glTexSubImage2D(GL_TEXTURE_RECTANGLE, 0, 0, 0, w, h, fmt, GL_UNSIGNED_BYTE, d)
        self._reset_unpack_alignment()

        self._updated = True

    def update(self):
        self._dirty = True

    def activate(self):
        if not self._created:
            self._create()

        if self._dirty:
            self._update()

        glBindTexture(GL_TEXTURE_RECTANGLE, self._texid)
        glTexParameteri(GL_TEXTURE_RECTANGLE, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_RECTANGLE, GL_TEXTURE_MIN_FILTER, GL_NEAREST)

    def deactivate(self):
        glBindTexture(GL_TEXTURE_RECTANGLE, 0)

    def _pre_draw(self, spaces: pxng.Spaces):
        self.activate()
        spaces.model.push()
        spaces.texture.push()

    def _post_draw(self, spaces: pxng.Spaces):
        self.deactivate()
        spaces.model.pop()
        spaces.texture.pop()

    def draw(self, spaces: pxng.Spaces):
        """
        Draw the sprite at the current position.

        Parameters
        ----------
        spaces : pxng.Spaces
            the current coordinate systems
        """
        self._pre_draw(spaces)
        spaces.model.translate((0, self._height, 0))  # anchor is upper left
        spaces.model.scale((self._width, self._height, 1))
        spaces.texture.scale(glm.vec3(self._width, self._height, 1))
        self._rect.draw(spaces)
        self._post_draw(spaces)

    def draw_partial(self, spaces: pxng.Spaces, x, y, width, height):
        """
        Draw a partial sprite using pixel coordinates for the partial image data.
        The coordinates represent a sub region in the sprite.

        Parameters
        ----------
        spaces: pxng.Spaces
            the current coordinate systems
        x: int
            x coordinate in pixel space
        y: int
            y coordinate in pixel space
        width: int
            width in number of pixels
        height: int
            height in number of pixels
        """
        self._pre_draw(spaces)
        spaces.model.translate((0, height, 0))  # anchor is upper left
        spaces.model.scale((width, height, 1))

        spaces.texture.translate(glm.vec3(x, y, 0))
        spaces.texture.scale(glm.vec3(width, height, 1))
        self._rect.draw(spaces)

        self._post_draw(spaces)

    @classmethod
    def create_from_image(cls, path):
        img_data = imageio.imread(path)
        return Sprite(img_data)


from typing import Tuple

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, confusion_matrix

class Evaluator:
    """
    Evaluator Class
    """

    def __init__(self):
        self.mapping = {
            0: "negative",
            1: "negative",
            2: "neutral",
            3: "neutral",
            4: "positive",
            5: "positive",
        }
        self.dataset = None
        self.train = None
        self.test = None
        self.model = None

    def summarize_review(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        calculates overall sentiment for each aspect of a single review
        using the polarity strength of each occurence of the aspect

        Args:
            data (pd.DataFrame): Dataframe containing information about
            the aspects, their polarities and their respective review

        Returns:
            pd.DataFrame: Dataframe with overall polarity for each aspect in a review
        """

        data["review_polarity"] = data.groupby(
            ["reviewnumber", "aspect"], as_index=False
        )["polarity_strength"].transform(lambda x: ",".join(x))
        data = data.drop(
            columns=[
                "sent_idx",
                "word_idx",
                "word_found",
                "polarity_strength",
                "sentiment_words",
                "intensifier_words",
            ]
        ).drop_duplicates()

        def calculate_review_polarity(review_polarity):
            review_polarity = review_polarity.replace("[", "")
            review_polarity = review_polarity.replace("]", "")
            review_polarity = review_polarity.split(",")
            try:
                review_polarity = np.mean(
                    [float(polarity) for polarity in review_polarity]
                )
            except ValueError:
                review_polarity = np.nan
            return review_polarity

        data["review_polarity"] = data["review_polarity"].apply(
            lambda x: calculate_review_polarity(x)
        )
        data = data.dropna()
        return data

    def read_data(
        self, path: str = "src/data/data_aspects_tokens.csv"
    ) -> Tuple[list, list]:
        """
        reads csv data from specified path

        Args:
            path (str, optional): Path to the csv file. Defaults to 'src/data/data_aspects_tokens.csv'.

        Returns:
            Tuple[list, list]: list with review polarity for an aspect and list with corresponding true label
        """
        data = pd.read_csv(path)
        data = self.summarize_review(data)
        self.dataset = data
        x = data["review_polarity"].to_list()
        y = data["true_label"].astype(int)

        return x, y

    def sample_data(self, x: list, y: list) -> Tuple[list, list, list, list]:
        """
        samples data so that possible labels are equally distributed

        Args:
            x (list): x-data (review polarity for aspect)
            y (list): y-data (true label for aspect in review)

        Returns:
            Tuple[list, list, list, list]: dataset splitted into train and test set
        """
        y = [self.mapping[i] for i in y]

        x_train = list()
        x_test = list()
        y_test = list()
        y_train = list()

        min_label = min([y.count(i) for i in list(set(y))])
        nbr_labels = len(set(y))
        data_size = min_label * len(set(y))

        train_size = int(0.66 * data_size)
        test_size = data_size - train_size

        for elem_x, elem_y in zip(x, y):
            if y_train.count(elem_y) < int(train_size / nbr_labels):
                x_train.append([elem_x])
                y_train.append(elem_y)

            elif y_test.count(elem_y) < int(test_size / nbr_labels):
                x_test.append([elem_x])
                y_test.append(elem_y)

        return x_train, y_train, x_test, y_test

    def generate_train_test(self):
        """
        generates train and test set
        """
        x, y = self.read_data()
        x_train, y_train, x_test, y_test = self.sample_data(x, y)
        self.train = (x_train, y_train)
        self.test = (x_test, y_test)

    def train_model(self):
        """
        trains model
        """
        model = LogisticRegression()
        model.fit(self.train[0], self.train[1])
        self.model = model

    def plot_results(self, predictions: list):
        """
        creates and saves a confusion matrix and boxplot

        Args:
            predictions (list): list containing the predicted labels for the test data
        """
        fig, ax = plt.subplots()
        cm = confusion_matrix(self.test[1], predictions)
        conf = confusion_matrix(self.test[1], predictions).ravel()
        nbr_labels = len(set(self.test[1]))
        cm = conf.reshape(nbr_labels, nbr_labels)
        sns.heatmap(cm, annot=True, fmt="d", cmap="Spectral")
        ax.set_xlabel("predicted label")
        ax.set_ylabel("true label")
        fig.savefig("confusion_matrix")

        fig, ax = plt.subplots()
        x = self.train[0] + self.test[0]
        y = self.train[1] + self.test[1]
        x = [i[0] for i in x]
        y = [i for i in y]
        results = pd.DataFrame({"polarity strength": x, "true label": y})
        sns.boxplot(data=results, x="true label", y="polarity strength")
        fig.savefig("boxplot")

    def evaluate(self):
        """
        performs evaluation of the results including visualization
        """
        predictions = self.model.predict(self.test[0])
        accuracy = accuracy_score(self.test[1], predictions)
        print("Accuracy:", str(accuracy * 100) + "%")
        self.plot_results(predictions)


#!/usr/bin/python3

from __future__ import print_function
import os.path
import pandas as pd
import numpy as np
from scipy.spatial import cKDTree

class Contact:
    """Parse a PDB/DSSP/CCP4 file and returns a DataFrame object."""
    def __init__(self, structure):
        """

        """
        self._structure = structure
        self.all2all = self.get_all2all()
        self.water2water = self.get_water2water()
        self.water2res = self.get_water2res()
        self.water2all = self.get_water2all()
        self.res2res = self.get_res2res()

        # try:
        #     self.get_water2water()
        #     self.get_water2res()
        #     self.get_res2res()
        # except:
        #     pass

    def get_all2all(self):
        df = self._structure.allatoms.copy()
        coords = np.array(df[['COOR_X','COOR_Y','COOR_Z']].values.astype(np.float))
        tree = cKDTree(coords)
        pairs = list(tree.query_pairs(r=4.9))

        idx1 = list(list(zip(*pairs))[0])
        idx2 = list(list(zip(*pairs))[1])

        df1 = df.iloc[idx1, :].copy().reset_index(drop=True).add_suffix('_1')
        df2 = df.iloc[idx2, :].copy().reset_index(drop=True).add_suffix('_2')
        dist = np.linalg.norm(coords[idx1]-coords[idx2], axis=1)

        all2all_df = pd.concat([df1, df2], axis=1)
        all2all_df['DIST'] = dist
        return all2all_df

    def get_water2water(self):
        df = self.all2all
        water2water_df = df.loc[(df.RES_NAME_1=='HOH')&\
                                 (df.ATOM_NAME_1=='O')&\
                                 (df.RES_NAME_2=='HOH')&\
                                 (df.ATOM_NAME_2=='O')&\
                                 (df.DIST<3.2)].copy().reset_index(drop=True)
        return water2water_df

    def get_water2res(self):
        df = self.all2all
        water2res_df = df.loc[(df.RES_NAME_1!=df.RES_NAME_2)&\
                              (df.ATOM_NAME_1!=df.ATOM_NAME_2)&\
                              (((df.RES_NAME_1=='HOH')&(df.ATOM_NAME_1=='O'))|\
                                (df.RES_NAME_2=='HOH')&(df.ATOM_NAME_2=='O'))&\
                              (df.DIST<3.2)].copy().reset_index(drop=True)
        return water2res_df

    def get_water2all(self):
        water2water_df = self.water2water.copy()
        water2res_df = self.water2res.copy()
        return pd.concat([water2water_df,water2res_df],axis=0)

    def get_res2res(self):

        df = self.all2all

        backbone_atoms = ['N', 'C', 'CA', 'O', 'CB', 'OXT']
        sidechain_atoms = ['SD', 'CZ', 'CH2', 'CG1', 'NZ', \
                           'ND2', 'OD1', 'CD1', 'CE1', 'CE', \
                           'CD2', 'CZ3', 'SG', 'OE2', 'NE1', \
                           'NE', 'CE2', 'CG', 'NH2', 'OE1', \
                           'CD', 'ND1', 'OH', 'NE2', 'NH1', \
                           'CG2', 'OD2', 'CZ2', 'CE3', 'OG', 'OG1']
        res_names = ['ALA', 'ARG', 'ASN', 'ASP', 'CYS', \
                     'GLN', 'GLU', 'GLY', 'HIS', 'ILE', \
                     'LEU', 'LYS', 'MET', 'PHE', 'PRO', \
                     'SER', 'THR', 'TRP', 'TYR', 'VAL']

        res2res_df = df.loc[(df.RES_NAME_1.isin(res_names))&\
                            (df.ATOM_NAME_1.isin(backbone_atoms+sidechain_atoms))&\
                            (df.RES_NAME_2.isin(res_names))&\
                            (df.ATOM_NAME_2.isin(backbone_atoms+sidechain_atoms))]

        return res2res_df

    def get_vector(self, atom1, atom2):
        structure = self._structure
        return structure.get_coordinate(atom2)-structure.get_coordinate(atom1)


"""
This creates Figure S7.
"""
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.cm as cm
from .figureCommon import subplotLabel, getSetup, plot_conf_int, import_pMuteins, organize_expr_pred, expScaleMut
from ..imports import import_Rexpr, import_samples_2_15, import_pstat

dataMean = import_pMuteins()
dataMean.reset_index(inplace=True)
data, _, _ = import_Rexpr()
data.reset_index(inplace=True)
unkVec_2_15 = import_samples_2_15(N=25)
_, _, _, _, pstat_df = import_pstat()
dataMean = dataMean.append(pstat_df, ignore_index=True, sort=True)

def makeFigure():
    """Get a list of the axis objects and create a figure"""
    # Get list of axis objects
    ax, f = getSetup((12, 18), (8, 6))

    subplotLabel(ax)

    tps = np.array([0.5, 1.0, 2.0, 4.0]) * 60.0
    tpsScale = np.array([0.5, 1.0]) * 60.0
    muteinC = dataMean.Concentration.unique()
    dataMean["Concentration"] = dataMean["Concentration"].astype(np.float)  # logscale for plotting

    ligand_order = ["WT N-term", "WT C-term", "V91K C-term", "R38Q N-term", "F42Q N-Term", "N88D C-term"]
    cell_order = ["NK", "CD8+", "T-reg", "Naive Treg", "Mem Treg", "T-helper", "Naive Th", "Mem Th"]

    df = pd.DataFrame(columns=["Cells", "Ligand", "Time Point", "Concentration", "Activity Type", "Replicate", "Activity"])  # make empty dataframe for all cell types

    # loop for each cell type and mutein
    for _, cell_name in enumerate(cell_order):
        IL2Ra = data.loc[(data["Cell Type"] == cell_name) & (data["Receptor"] == "IL-2R$\\alpha$"), "Count"].values[0]
        IL2Rb = data.loc[(data["Cell Type"] == cell_name) & (data["Receptor"] == "IL-2R$\\beta$"), "Count"].values[0]
        gc = data.loc[(data["Cell Type"] == cell_name) & (data["Receptor"] == "$\\gamma_{c}$"), "Count"].values[0]
        receptors = np.array([IL2Ra, IL2Rb, gc]).astype(np.float)

        for _, ligand_name in enumerate(ligand_order):
            # append dataframe with experimental and predicted activity
            df = organize_expr_pred(df, cell_name, ligand_name, receptors, muteinC, tps, unkVec_2_15)

    df = expScaleMut(df, tpsScale)

    # determine scaling constants
    plot_expr_predM(ax, df, cell_order, ligand_order, tps, muteinC)

    return f

def plot_expr_predM(ax, df, cell_order, ligand_order, tps, muteinC):
    """ Plots experimental and scaled model-predicted dose response for all cell types, muteins, and time points. """
    pred_data = np.zeros((12, 4, unkVec_2_15.shape[1]))
    bounds = np.array([800, 50, 150, 150, 150, 15, 15, 15])

    for i, cell_name in enumerate(cell_order):
        for j, ligand_name in enumerate(ligand_order):
            axis = i * 6 + j

            # plot experimental data
            sns.scatterplot(x="Concentration", y="Activity", hue="Time Point", data=df.loc[(df["Cells"] == cell_name) & (
                df["Ligand"] == ligand_name) & (df["Activity Type"] == "experimental")], ax=ax[axis], s=10, palette=cm.rainbow, legend=False)

            # scale and plot model predictions
            for k, conc in enumerate(df.Concentration.unique()):
                for l, tp in enumerate(tps):
                    for m in range(unkVec_2_15.shape[1]):
                        pred_data[k, l, m] = df.loc[
                            (df["Cells"] == cell_name)
                            & (df["Ligand"] == ligand_name)
                            & (df["Activity Type"] == "predicted")
                            & (df["Concentration"] == conc)
                            & (df["Time Point"] == tp)
                            & (df["Replicate"] == (m + 1)),
                            "Activity",
                        ]

            plot_dose_responseM(ax[axis], pred_data, tps, muteinC, legend=(axis == 0))
            ax[axis].set(xlabel=("[" + ligand_name + "] (nM)"), ylabel="Activity", title=cell_name, ylim=(0, bounds[i]))

def plot_dose_responseM(ax, mutein_activity, tps, muteinC, legend=False):
    """ Plots predicted activity for multiple timepoints and mutein concentrations. """
    colors = cm.rainbow(np.linspace(0, 1, tps.size))

    for tt in range(tps.size):
        if not legend:
            plot_conf_int(ax, muteinC.astype(np.float), mutein_activity[:, tt, :], colors[tt])
        else:
            plot_conf_int(ax, muteinC.astype(np.float), mutein_activity[:, tt, :], colors[tt], (tps[tt] / 60.0).astype(str))
            ax.legend(title="time (hours)")
    ax.set_xscale("log")
    ax.set_xticks(np.array([10e-5, 10e-3, 10e-1, 10e1]))


import enum

import cv2
import numpy
import requests

import config
from model.model import KerasModel
from model.training_data import TrainingData

class EmoteType(enum.Enum):
    PEEPO = 0
    KAPPA = 1

class EmoteRecognizer:

    def __init__(self):
        training_data = TrainingData(config.peepo_data_dir, config.kappa_data_dir, config.image_size)

        keras_model = KerasModel(config.model_filepath, config.image_size, config.batch_size, training_data)
        keras_model.create_weights()

        self.model = keras_model.model
        self.image_size = config.image_size

    def predict(self, image_array):
        """
        Predicts an image

        :param image_array: the image in the correct size turned into an array
        :return: the predicted emote type of the image
        """

        # the model expects a list a images to predict, but we have only one image.
        # So we're expanding the array
        image_array = numpy.expand_dims(image_array, 0)

        prediction = self.model.predict(image_array)[0][0]
        return EmoteType(prediction)

    def parse_image(self, url):
        """
        Downloads an image, resizes it and turns it into an array

        :param url: the url of the image which should be downloaded
        :return: the array of the image, ready to be predicted
        """

        response = requests.get(url)
        image = numpy.asarray(bytearray(response.content), dtype="uint8")
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)
        return cv2.resize(image, (self.image_size, self.image_size), interpolation=cv2.INTER_CUBIC)


import re
from serial import Serial
import time
import numpy as np
import ML.FeatureExtractor as features
import threading
import Queue
import camera as mycamera
import os
import copy

#queue = Queue.Queue()
#camera_done = False

def get_height_and_width(data):
    measures = []
    measures.append(213.5 - float(data[0]))
    measures.append(112 - float(data[1]) - float(data[2]))
    return measures

def write_event_data(event_file, event_data, len, current_time):
    for i in range(0,len):
        event_file.write(str(event_data[i,0])+','+str(event_data[i,1])+'\n')
    event_file.write(str(current_time) + '_\n')
    event_file.flush()
"""
def get_video_events(video_id):
    print 'video_id: ' , video_id
    while not queue.empty():
        info = queue.get()
        print 'queue item: ', info
        if info[3] != video_id:
            queue.put(info)
            break
        video_time = info[4]
        my_start = info[0] - video_time - 1
        if my_start < 0:
            my_start = 0
        my_end = info[1] - video_time + 1
        if my_end > 10:
            my_end = 10
        event_id = info[2]
        mycamera.extract_video_event(my_start, my_end, event_id, video_id)
    #delete video
    print 'deleting ' + 'output' + str(video_id) + '.avi'
    os.remove('output'+str(video_id)+'.avi')

#A wrapper around the do_capture function in Camera module
def do_capture(id):
    global camera_done
    camera_done = False
    mycamera.do_capture(id)
    camera_done = True

def get_video_event():
    #Camera footage stuff
    video_time = time.time()
    video_id = 1
    camera_thread = threading.Thread(target=do_capture,args=(video_id,))
    #camera_thread.daemon = True
    camera_thread.start()
    video_id += 1
"""
def read_sensor_data():
    ser = Serial('/dev/ttyACM0',9600)
    count= 0
    temperature = 23
    calibration_a = 0.013696
    calibration_b = -3.56201
    event = []
    eventCount = 0
    event_id = num_lines = sum(1 for line in open('features.dat')) + 1
    all_data_file = open('all_data.dat','a')
    raw_event_file = open('event_raw.data', 'a')
    feature_file = open('features.dat','a')
    log_file = open('readline.log','a')
    event_time = []

    #Camera footage stuff
    #video_time = time.time()
    #video_id = 1
    #camera_thread = threading.Thread(target=do_capture,args=(int(video_id),))
    #camera_thread.daemon = True
    #camera_thread.start()
    #start reading sensory data
    while True:
        try:
            line = ser.readline()
            current_time = time.time()
            count +=1
            data = line.split(',')
            """
            #camera footage stuff
            if camera_done: # this means video shot finished

                #Extracting video from 1-minute long video
                subvideo_thread = threading.Thread(target=get_video_events,args=(int(video_id),))
                #subvideo_thread.daemon = True
                subvideo_thread.start()
                #Start new video
                video_id += 1
                video_time = time.time()
                camera_thread = threading.Thread(target=do_capture,args=(int(video_id),))
                #camera_thread.daemon = True
                camera_thread.start()
            """

            all_data_file.write(str(current_time)+','+line+'\n')
            all_data_file.flush()
            if len(data) == 3:
                for i in range(0,3):
                    val =re.sub('[^0-9]','', data[i])
                    #data[i] = int(val)#((float(val) / 10**4 ) * (331.3 + 0.606 * temperature)) / 2
                    data[i] = calibration_a * float(val) + calibration_b

                measures = get_height_and_width(data)

                if (110 < measures[0] and measures[0] < 200) or measures[1] > 20:  #start of passing event
                    #print "measures: ",measures
                    if eventCount == 0:
                        event_time = []
                        event_time.append(copy.copy(time.time()))
                    event.append(measures)
                    eventCount += 1
                elif eventCount > 5: # extract features
                    event_time.append(time.time())
                    eventdata = np.array(event)
                    avgHeight = features.getAvgHeight(eventdata)
                    maxHeight = features.getMaxHeight(eventdata)
                    avgWidth = features.getAvgWidth(eventdata)
                    maxWidth = features.getMaxWidth(eventdata)
                    circumference = features.getcircumference(eventdata, 50.0)
                    feature_line =  str(avgHeight)+','+ str(maxHeight)+','+ str(avgWidth)+ ','+ str(maxWidth) + ',' + str(circumference)
                    #print feature_line
                    feature_file.write(str(current_time)+','+feature_line+'\n')
                    feature_file.flush()
                    write_event_data(raw_event_file,eventdata,eventCount,current_time)

                    #save event id to structrure and increment event_id
                    event_time.append(copy.copy(event_id))
                    #event_time.append(copy.copy(video_id))
                    #event_time.append(copy.copy(video_time))
                    event_id += 1
                    #push event time to queue
                    #print 'pushing ', event_time
                    #queue.put(event_time)
                    #print 'queue size: ', queue._qsize()
                    #initialize vars for next event
                    eventCount = 0
                    event = []
                    event_time = []
        except Exception:
            t = time.time()
            print 'error'
            log_file.write('An error occured at ' + str(t))
            pass

#start program
read_sensor_data()