#!/usr/bin/env python
import os
from merger import Merger
from unimerger import UniMerger
from alignerbwa import AlignerBwa
from alignerbbmap import AlignerBBMap
from counterfc import CounterFC
from samfc import SamFC
from counterjl import CounterJL
from metrics import Metrics
from subprocess import call
from alignerutils import sort_by_qname
from alignerutils import sort_bam

# This would execute the core part for RtS (RNATag-Seq) pipeline.

class RtSCore:
    def __init__(self, cldict, sampd):
        self.cldict = cldict
        self.sampd = sampd
        self.mergo = Merger(cldict, sampd)
        self.meto = Metrics(cldict)

        lbwao = None
        lbbmapo = None
        lref_acc_str = sampd.ref_acc_str
        lhost_ref_str = sampd.host_ref_str

        if lref_acc_str != "none":
            lbwao = AlignerBwa(cldict, sampd)
        if lhost_ref_str != "none":
            lbbmapo = AlignerBBMap(cldict, sampd)
        self.bwao = lbwao
        self.bbmapo = lbbmapo
        self.samfco = SamFC(cldict, sampd)
        self.jlco = CounterJL(cldict, sampd)
        print("Created JLCounter object")

    def merged_single_name(self, sample_id, suffix, merge_dir, ldelim):
        outfile_R = merge_dir + ldelim + sample_id + "_" + suffix + ".fastq"
        cldict = self.cldict
        gzip_merged = cldict.gzip_merged
        if gzip_merged:
            outfile_R += ".gz"
        elif not os.path.isfile(outfile_R):
            outfile_R += ".gz"

        if not os.path.isfile(outfile_R):
            raise StandardError("Merged file does not exist: " + outfile_R)

        return outfile_R

    def fastq_to_bam_single(self):
        cldict = self.cldict
        sampd = self.sampd
        mergo = self.mergo
        bwao = self.bwao

        ldelim = cldict.ldelim
        Patho_dir = cldict.Patho_dir
        merge_dir = cldict.Merge_dir
        project_id = sampd.project_id
        Bam_path = cldict.Bam_path
        bamdir = Bam_path + ldelim + project_id
        outdir = Patho_dir + ldelim + project_id
        sample_id = sampd.sample_id
        ref_acc = sampd.ref_acc_str

        suffix = "R"

        print("sample_id: " + sample_id)
        print("outdir: " + outdir)
        read_file = self.merged_single_name(sample_id, suffix, merge_dir, ldelim)
        sorted_bam = bwao.exe_bwa_single(sample_id, ref_acc, read_file, suffix, outdir, bamdir)
        return sorted_bam

    def fastq_to_bam_paired(self):
        cldict = self.cldict
        sampd = self.sampd
        mergo = self.mergo
        bwao = self.bwao

        ldelim = cldict.ldelim
        Patho_dir = cldict.Patho_dir
        project_id = sampd.project_id
        Bam_path = cldict.Bam_path
        bamdir = Bam_path + ldelim + project_id
        outdir = Patho_dir + ldelim + project_id
        sample_id = sampd.sample_id
        ref_acc = sampd.ref_acc_str
        merge_dir = cldict.Merge_dir

        suf1 = "R1"
        suf2 = "R2"

        print("sample_id: " + sample_id)
        print("outdir: " + outdir)
        read1_file = self.merged_single_name(sample_id, suf1, merge_dir, ldelim)
        read2_file = self.merged_single_name(sample_id, suf2, merge_dir, ldelim)
        sorted_bam = bwao.exe_bwa_paired(sample_id, ref_acc, read1_file, read2_file,
            suf1, suf2, outdir, bamdir)
        return sorted_bam

    def sample_id_to_bam_patho(self):
        cldict = self.cldict
        sampd = self.sampd
        Patho_dir = cldict.Patho_dir
        ref_acc = sampd.ref_acc_str
        ldelim = cldict.ldelim
        project_id = sampd.project_id
        outdir = Patho_dir + ldelim + project_id
        sample_id = sampd.sample_id
        pair_suffix = None
        Read_pairing_val = cldict.Read_pairing_val
        if Read_pairing_val == 'SINGLE':
            pair_suffix = 'se'
        elif Read_pairing_val == 'PAIRED':
            pair_suffix = 'pe'
        outbamfile = outdir + ldelim + sample_id + "_" + ref_acc + "_" + \
            pair_suffix + ".bam"
        return outbamfile

    def exe_patho(self):
        cldict = self.cldict
        meto = self.meto
        bwao = self.bwao
        sampd = self.sampd
        Patho_dir = cldict.Patho_dir
        ldelim = cldict.ldelim
        ldel = cldict.ldelim
        project_id = sampd.project_id
        samtools_path = cldict.samtools
        tdf_path = cldict.tdf_str

        Bam_path = cldict.Bam_path
        bamdir = Bam_path + ldelim + project_id

        outdir = Patho_dir + ldelim + project_id
        Results_path = cldict.Results_path
        picard_outdir = Results_path + ldel + project_id + ldel + "bam_metrics_patho"
        nodup_picard_outdir = Results_path + ldel + project_id + ldel + "nodupdir" + ldel + "bam_metrics_patho"
        sample_id = sampd.sample_id
        ref_acc = sampd.ref_acc_str
        Read_pairing_val = cldict.Read_pairing_val
        rm_rts_dup = cldict.rm_rts_dup
        jlco = self.jlco

        sorted_bam = ''
        do_align = cldict.do_align
        if do_align:
            if Read_pairing_val == 'SINGLE':
                sorted_bam = self.fastq_to_bam_single()
            elif Read_pairing_val == 'PAIRED':
                print("Started fastq_to_bam_paired from exe_patho in RtSCore")
                sorted_bam = self.fastq_to_bam_paired()
            else:
                raise StandardError("Unknown Read_pairing: " + Read_pairing_val)

        else:
            sorted_bam = self.sample_id_to_bam_patho()

        do_count = cldict.do_count
        if do_count:
            read_counter = cldict.read_counter
            if read_counter == 'featureCounts':
                fco = CounterFC(cldict, sampd)
                print ("Created featureCount object")
                countfile_s_str = fco.exe_featureCounts(sample_id, ref_acc, sorted_bam, outdir, "s")
                countfile_as_str = fco.exe_featureCounts(sample_id, ref_acc, sorted_bam, outdir, "as")
                countfile_str = outdir + ldelim + sample_id + "_" + ref_acc + ".counts"
                CounterFC.combine_s_as(countfile_s_str, countfile_as_str, countfile_str)
            elif read_counter == 'JL_counter':
                if Read_pairing_val == 'SINGLE':
                    jlco.count_single(sample_id, ref_acc, sorted_bam, outdir, strand_rev = "N")
                elif Read_pairing_val == 'PAIRED':
                    jlco.count_paired(sample_id, ref_acc, sorted_bam, outdir)
                fnafile = bwao.get_patho_ref_path(ref_acc)
                meto.exe_picard_metrics(sorted_bam, fnafile, picard_outdir)
            else:
                raise StandardError('No read counter found for rts pathogen side.')

            if rm_rts_dup:
                # Do the duplicate removal to a new directory called no_dup_dir
                no_dup_dir = outdir + ldelim + "nodupdir"
                # Create a new bam file in the no_dup_dir after removal of the
                # pcr duplicates from pcr; ideally that folder should also
                # contain the metrics related to dup removal
                dup_marked_sorted_bam = self.mark_dup_reads(sorted_bam)
                nodup_sorted_bam = self.remove_dup_reads(dup_marked_sorted_bam)
                if Read_pairing_val == 'SINGLE':
                    jlco.count_single(sample_id, ref_acc, nodup_sorted_bam, no_dup_dir, strand_rev = "N")
                elif Read_pairing_val == 'PAIRED':
                    jlco.count_paired(sample_id, ref_acc, nodup_sorted_bam, no_dup_dir)
                fnafile = bwao.get_patho_ref_path(ref_acc)
                meto.exe_picard_metrics(nodup_sorted_bam, fnafile, nodup_picard_outdir)
                # We have to copy the bam file to appropriate place
                nodup_bamdir = bamdir + ldelim + "nodupdir"
                AlignerBwa.copy_bam_tdf(samtools_path, tdf_path, nodup_bamdir, nodup_sorted_bam)

    def remove_dup_reads(self, dm_sorted_bam):
        cldict = self.cldict
        samtools_path = cldict.samtools
        remove_dup_script = cldict.remove_dup_script

        # 1. sort wrt queryname
        dm_unsorted_bam = dm_sorted_bam.replace("_dm.bam", "_u_dm.bam")
        if not os.path.isfile(dm_unsorted_bam):
            sort_by_qname(samtools_path, dm_sorted_bam, dm_unsorted_bam)
            print("sort by qname: " + dm_sorted_bam + " to " + dm_unsorted_bam)

        # 2. remove the pair of marked reads 0x400
        nodup_unsorted_bam = dm_unsorted_bam.replace("_u_dm.bam", "_u.bam")
        dup_bam = dm_unsorted_bam.replace("_u_dm.bam", "_u_dup.bam")
        # Call the c++ program remove_dup
        remove_dup_cmd = remove_dup_script + " -i " + dm_unsorted_bam +\
            " -o " + nodup_unsorted_bam + " -d " + dup_bam
        print("remove_dup_cmd starts: " + remove_dup_cmd)
        call(remove_dup_cmd.split())
        print("remove_dup_cmd ends")

        # 3. sort the no_dup_unsorted_bam bam file
        nodup_sorted_bam = nodup_unsorted_bam.replace("_u.bam", "_pe.bam")
        if not os.path.isfile(nodup_sorted_bam):
            sort_bam(samtools_path, nodup_unsorted_bam, nodup_sorted_bam)
            print("sort by coordinate: " + nodup_unsorted_bam + " to " +\
                nodup_sorted_bam)
        os.remove(dm_unsorted_bam)
        print("Removed dm_unsorted_bam: " + dm_unsorted_bam)
        os.remove(dm_sorted_bam)
        print("Remove dm_sorted_bam: " + dm_sorted_bam)
        return nodup_sorted_bam

    def mark_dup_reads(self, input_bam):
        # Get the output directory which has been created before
        # Do not remove the duplicates, just let picard mark them using read is PCR or optical
        # duplicate (0x400). We shall remove them later in pair
        cldict = self.cldict
        ldelim = cldict.ldelim
        picard_bindir = cldict.picard_bindir
        picard_jar = picard_bindir + "/picard.jar"

        bamdir = os.path.dirname(input_bam)
        no_dup_dir = bamdir + ldelim + "nodupdir"
        bam_file = os.path.basename(input_bam)
        output_bam_ori = no_dup_dir + ldelim + bam_file
        # Output bam is the dup marked one
        output_bam = output_bam_ori.replace("_pe.bam", "_dm.bam")

        metrics_txt = output_bam.replace('.bam', '_rm_dup_met.txt')

        # Then rename the old bam file

        mem_str = "-Xmx6G"
        dup_str = "MarkDuplicates"
        picard_cmd = "java " + mem_str + " -jar " + picard_jar + " " + dup_str + " I=" + input_bam +\
            " O=" + output_bam + " M=" + metrics_txt + " REMOVE_DUPLICATES=false VALIDATION_STRINGENCY=LENIENT"
        print("start: picard_cmd: " + picard_cmd)
        call(picard_cmd.split())
        print("end: picard_cmd:")
        return (output_bam)

    def sample_id_to_unsorted_bam_path(self):
        sampd = self.sampd
        lhost_ref_str = sampd.host_ref_str
        cldict = self.cldict
        ldelim = cldict.ldelim
        Host_dir = cldict.Host_dir
        project_id = sampd.project_id
        outdir = Host_dir + ldelim + project_id
        sample_id = sampd.sample_id
        unsorted_bam_path = outdir + ldelim + sample_id + "_" + lhost_ref_str + "_u.bam"
        if not os.path.exists(unsorted_bam_path):
            raise StandardError("File not found: " + unsorted_bam_path)
        return unsorted_bam_path

    def sample_id_to_sorted_bam_path(self, pair_type):
        sampd = self.sampd
        lhost_ref_str = sampd.host_ref_str
        cldict = self.cldict
        ldelim = cldict.ldelim
        Host_dir = cldict.Host_dir
        project_id = sampd.project_id
        outdir = Host_dir + ldelim + project_id
        sample_id = sampd.sample_id
        unsorted_bam_path = outdir + ldelim + sample_id + "_" + lhost_ref_str + "_" + pair_type + ".bam"
        if not os.path.exists(unsorted_bam_path):
            raise StandardError("File not found: " + unsorted_bam_path)
        return unsorted_bam_path

    def host_fastq_to_bam_paired(self):
        """
        """
        sampd = self.sampd
        cldict = self.cldict
        Host_dir = cldict.Host_dir
        ldelim = cldict.ldelim
        project_id = sampd.project_id
        outdir = Host_dir + ldelim + project_id
        Bam_path = cldict.Bam_path
        bamdir = Bam_path + ldelim + project_id
        sample_id = sampd.sample_id
        bbmapo = self.bbmapo
        merge_dir = cldict.Merge_dir
        do_align = cldict.do_align
        do_count = cldict.do_count

        suf1 = "R1"
        suf2 = "R2"

        print("sample_id: " + sample_id)
        print("outdir: " + outdir)
        read1_file = self.merged_single_name(sample_id, suf1, merge_dir, ldelim)
        read2_file = self.merged_single_name(sample_id, suf2, merge_dir, ldelim)
        if do_align:
            bbmapo.exe_host_paired(sample_id, read1_file, read2_file, outdir, bamdir)

        #bbmapo.exe_fragcount(sample_id, outdir)

    def host_bam_to_count(self, unsorted_bam):
        """
        """
        sampd = self.sampd
        sample_id = sampd.sample_id
        cldict = self.cldict
        Host_dir = cldict.Host_dir
        ldelim = cldict.ldelim
        project_id = sampd.project_id
        outdir = os.path.dirname(unsorted_bam)
        samfco = self.samfco
        proto = "rts"
        end = "paired"
        out_count, basic_metrics = samfco.exe_count_frags(sample_id, unsorted_bam, outdir, proto, end)
        out_gene_count = samfco.exe_frag_to_gene_count(sample_id, out_count, outdir)
        detailed_metrics_file = samfco.exe_detailed_metrics(sample_id, out_gene_count, basic_metrics, outdir)

    def get_unsorted_bam(self, sorted_bam):
        # Check if the unsorted bam exists and if not just generates it.
        cldict = self.cldict
        samtools_path = cldict.samtools
        unsorted_bam = None

        if not os.path.isfile(sorted_bam):
            lstr = "Bam file not found: " + sorted_bam
            raise StandardError(lstr)

        if sorted_bam.endswith("_pe.bam"):
            unsorted_bam = sorted_bam.replace("_pe.bam", "_u.bam")
        elif sorted_bam.endswith("_se.bam"):
            unsorted_bam = sorted_bam.replace("_se.bam", "_u.bam")

        if not os.path.isfile(unsorted_bam):
            sort_by_qname(samtools_path, sorted_bam, unsorted_bam)

        return (unsorted_bam)

    def exe_host(self):
        cldict = self.cldict
        sampd = self.sampd
        do_count = cldict.do_count
        Results_path = cldict.Results_path
        ldel = cldict.ldelim
        project_id = sampd.project_id
        ref_acc = sampd.ref_acc_str
        bbmapo = self.bbmapo
        rm_rts_dup = cldict.rm_rts_dup
        samtools_path = cldict.samtools
        tdf_path = cldict.tdf_str

        Bam_path = cldict.Bam_path
        bamdir = Bam_path + ldel + project_id

        Read_pairing_val = cldict.Read_pairing_val
        sorted_bam = None
        if Read_pairing_val == 'SINGLE':
            self.host_fastq_to_bam_single()
            sorted_bam = self.sample_id_to_sorted_bam_path("se")
        elif Read_pairing_val == 'PAIRED':
            self.host_fastq_to_bam_paired()
            sorted_bam = self.sample_id_to_sorted_bam_path("pe")
        else:
            raise ValueError('Incorrect value for Read_pairing_val: ' + Read_pairing_val)

        host_fna_path = bbmapo.get_host_fna_path()
        print("host_fna_path: " + host_fna_path)

        if do_count:
            unsorted_bam = self.sample_id_to_unsorted_bam_path()
            count_file = self.host_bam_to_count(unsorted_bam)
            picard_outdir = Results_path + ldel + project_id + ldel + \
                "bam_metrics_host"
            meto = self.meto
            print("picard_outdir: " + picard_outdir)
            meto.exe_picard_metrics(sorted_bam, host_fna_path, picard_outdir)

            if rm_rts_dup:
                outdir = Results_path + ldel + project_id
                # Do the duplicate removal to a new directory called no_dup_dir
                no_dup_dir = outdir + ldel + "nodupdir"
                # Create a new bam file in the no_dup_dir after removal of the pcr duplicates from pcr;
                # ideally that folder should also contain the metrics related to dup removal
                jlco = self.jlco
                dup_marked_sorted_bam = self.mark_dup_reads(sorted_bam)
                nodup_sorted_bam = self.remove_dup_reads(dup_marked_sorted_bam)

                nodup_unsorted_bam = self.get_unsorted_bam(nodup_sorted_bam)
                count_file = self.host_bam_to_count(nodup_unsorted_bam)
                nodup_picard_outdir = no_dup_dir + ldel + "bam_metrics_host"
                meto.exe_picard_metrics(nodup_sorted_bam, host_fna_path, nodup_picard_outdir)
                nodup_bamdir = bamdir + ldel + "nodupdir"
                AlignerBwa.copy_bam_tdf(samtools_path, tdf_path, nodup_bamdir, nodup_sorted_bam)

    def mainFunc(self):
        cldict = self.cldict
        do_host = cldict.do_host
        do_patho = cldict.do_patho
        sampd = self.sampd
        ref_acc_str = sampd.ref_acc_str
        ref_acc_str_l = ref_acc_str.lower()
        host_ref_str = sampd.host_ref_str
        host_ref_str_l = host_ref_str.lower()

        if do_patho:
            if ref_acc_str_l and ref_acc_str_l != "unknown" and ref_acc_str_l != "none":
                print("Started exe_patho from RtSCore")
                self.exe_patho()
        if do_host:
            if host_ref_str_l and host_ref_str_l != "unknown" and  host_ref_str_l != "none":
                self.exe_host()


import matplotlib.pyplot as plt
from sklearn import metrics
import seaborn as sns
from sklearn.neighbors import KNeighborsClassifier

#A python file that contains various functions for single use.

def plot_knn_f1s(X_train, X_test, y_train, y_test):
    k_scores=[]
    f_scores = []
    k_range = list(range(1, 20))
    for k in k_range:
        knn = KNeighborsClassifier(n_neighbors=k)
        knn.fit(X_train, y_train)
        y_pred=knn.predict(X_test)
        acc=metrics.accuracy_score(y_test,y_pred)
        f1 = metrics.f1_score(y_test,y_pred)
        k_scores.append(acc)
        f_scores.append(f1)
    list(zip(k_scores, f_scores))
    sns.set_style("darkgrid")
    plt.figure(figsize=(12, 6))
    plt.plot(k_range, f_scores, color='red', linestyle='dashed', marker='o',
             markerfacecolor='blue', markersize=10)
    plt.title('F1 score by K Value')
    plt.xlabel('K Value')
    plt.ylabel('Accuracy Score')
    plt.show()


import json
from utils.logger import log
from channels.consumer import AsyncConsumer
from channels.db import database_sync_to_async
from server.models import Object

class MapConsumer(AsyncConsumer):
    async def websocket_connect(self, event):
        log.info(event)

        self.chat_room = 'public'
        await self.channel_layer.group_add(
            self.chat_room,
            self.channel_name
        )
        log.debug("add channel layer")

        await self.send({
            "type": "websocket.accept"
        })
        log.debug("websocket.accept")

        await self.send({
            "type": "websocket.send",
            "text": json.dumps(
                {
                    "data": list(Object.objects.all().values('lat', 'lng', 'symbolCode'))
                })
        })
        log.debug(f"websocket.send data from db, {Object.objects.count()} object(s)")

    async def websocket_receive(self, event):
        log.info(event['type'])

        text = event.get('text', None)
        if text is not None:
            json_data = json.loads(text)
            log.debug(json_data['type'] + ' object')
            if json_data['type'] == 'add':
                await self.add_object(json_data['symbol'])
            elif json_data['type'] == 'delete':
                await self.del_object(json_data['symbol'])
            elif json_data['type'] == 'position':
                await self.move_object(json_data['symbol'])

            response = {"data": list([json_data])}

            await self.channel_layer.group_send(
                self.chat_room,
                {
                    "type": "chat_message",
                    "text": response
                }
            )

    async def chat_message(self, event):

        await self.send({
            "type": "websocket.send",
            "text": json.dumps(event['text'])
        })

    async def websocket_disconnect(self, event):
        log.info(event)

    @database_sync_to_async
    def add_object(self, text):
        try:
            Object.objects.create(lat=text['lat'], lng=text['lng'], symbolCode=text['symbolCode'])
            log.info("add to db")
        except Exception as e:
            log.error(e)

    @database_sync_to_async
    def del_object(self, text):
        try:
            Object.objects.get(lat=text['lat'], lng=text['lng'], symbolCode=text['symbolCode']).delete()
            log.info("delete from db: ")
        except Exception as e:
            log.error(e)

    @database_sync_to_async
    def move_object(self, text):
        try:
            o = Object.objects.get(lat=text['lat'], lng=text['lng'])
            o.lat = text['newLat']
            o.lng = text['newLng']
            o.save()
            log.info("move in db: ")
        except Exception as e:
            log.error(e)


# _*_ coding: utf-8 _*_

from scipy.optimize import leastsq
import numpy as np

vols=np.array([13.71,14.82,16.0,17.23,18.52])
energies=np.array([-56.29,-56.41,-56.46,-56.463,-56.41])

def Murnaghan(parameters,vol):
    'From Phys.Rev.B 28, 5480 (1983)'
    E0,B0,BP,V0=parameters
    E=E0+B0*vol/BP(((V0/vol)*BP)/(BP-1)+1)-V0*B0/(BP-1.0)
    return E

def objective(pars,y,x):
    # we will minimize this function
    err=y-Murnaghan(pars,x)
    return err

x0=[-56.0,0.54,2.0,16.5] #initial guess of parameters

plsq=leastsq(objective,x0,args=(energies,vols))

print('Fitted parameters = {0}'.format(plsq[0]))

# draw graph

import matplotlib.pyplot as plt

plt.plot(vols,energies,'ro')

# plot the fitted curve on top

x=np.linspace(min(vols),max(vols),50)
y=Murnaghan(plsq[0],x)
plt.plot(x,y,'k-')
plt.xlabel('Volume')
plt.ylabel('Energy')
plt.savefig('59.jpg',dpi=300)
plt.show()
print('plot done')


Mod = 1000000007

def fastpow(x,y,MOD):
    ans = 1
    while y:
        if (y & 1): ans = (ans * x) % MOD
        x = (x * x) % MOD
        y >>= 1
    return ans

s = input()
n = len(s)

cnt = [0 for _ in range(26)]
for x in s:
    cnt[ord(x) - ord('a')] += 1

fac = [1 for _ in range(n + 1)]
ifac = [1 for _ in range(n + 1)]
for i in range(1,n + 1):
    fac[i] = (fac[i - 1]*i) % Mod

ifac[n] = fastpow(fac[n],Mod - 2,Mod)
for i in range(n - 1,-1,-1):
    ifac[i] = (ifac[i + 1]*(i + 1)) % Mod

res = fac[n]
for i in range(26):
    res = (res*ifac[cnt[i]]) % Mod

print(res)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import mimetypes
import os
import re
import shutil
import subprocess

from bson import ObjectId
from pymongo import MongoClient

from google_takeout_util import (
    image_formats,
    video_formats,
    move_subfolders_into_batches,
)

def canonical_name(fn):
    bn = os.path.basename(os.path.splitext(fn)[0])
    bn = bn.replace("-edited", "")
    bn = re.sub(r"#\d+", "", bn)
    bn = re.sub(r"\(\d+\)", "", bn)
    bn = re.sub(r"[-_]\d$", "", bn)
    return bn.strip()

def _find_original_record(edited_record):
    # './Google Photos/2015-04-22-23/IMG_2079-edited.jpg' --> './Google Photos/2015-04-22-23/IMG_2079.JPG'
    original_base = (
        edited_record["SourceFile"].replace("-edited", "").replace("-redigerad", "")
    )
    myregexp = f".*{os.path.splitext(original_base)[0]}\..*"
    # print(myregexp)
    original_record = db.flatjson.find_one(
        {
            "$and": [
                {"SourceFile": {"$regex": myregexp}},
                {"MIMEType": {"$eq": edited_record["MIMEType"]}},
            ]
        }
    )
    return original_record

def move_edited_to_original():
    edited = media_collection.find(
        {
            "$or": [
                {"SourceFile": {"$regex": ".*-redigerad.*"}},
                {"SourceFile": {"$regex": ".*-edited.*"}},
            ]
        }
    )
    for e in edited:
        o = _find_original_record(e)
        if o:
            print(f"{e['SourceFile']} --> {os.path.basename(o['SourceFile'])}")
            try:
                shutil.move(
                    os.path.abspath(e["SourceFile"]), os.path.abspath(o["SourceFile"])
                )
            except FileNotFoundError:
                # assume the move happened already in an earlier run
                # print("J")
                pass
            db.media.delete_one({"_id": ObjectId(e["_id"])})

def fix_wrong_heic():
    wrong = media_collection.find(
        {
            "$and": [
                {"FileName": {"$regex": ".*\HEIC$"}},
                {"MIMEType": {"$ne": "image/heic"}},
            ]
        }
    )
    for w in wrong:
        new_ext = ""
        if w["MIMEType"] == "image/jpeg":
            new_ext = ".jpg"
        elif w["MIMEType"] == "video/mp4":
            new_ext = ".mp4"
        elif w["MIMEType"] == "video/quicktime":
            new_ext = ".mov"
        if not new_ext:
            print("-" * 50)
            print(w)
            print("-" * 50)
            continue
        change_extension(w, new_ext)

def change_extension(doc, new_ext):
    base, ext = os.path.splitext(doc["SourceFile"])
    mt = mimetypes.MimeTypes().guess_type(doc["SourceFile"])[0]
    print(mt)
    print(f"move {doc['SourceFile']} {base}.{new_ext}")
    try:
        shutil.move(doc["SourceFile"], f"{base}{new_ext}")
        print("ok")
    except FileNotFoundError:
        # assume the move happened already in an earlier run
        pass
    db.media.update_one(
        {"_id": doc["_id"]},
        {
            "$set": {
                "SourceFile": doc["SourceFile"].replace(ext, new_ext),
                "FileName": doc["FileName"].replace(ext, new_ext),
                "MIMEType": mt,
            }
        },
        upsert=False,
    )

def update_exif_date(doc, dt):
    command = [
        "/usr/local/bin/exiftool",
        "-m",
        "-overwrite_original",
        f"-DateTimeOriginal='{dt}'",
        f"-CreateDate='{dt}'",
        doc["SourceFile"],
    ]
    try:
        s = subprocess.check_output(command, encoding="utf-8", stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as e:
        if "Not a valid PNG (looks more like a JPEG)":
            change_extension(doc, ".jpg")
        print(e.output)

def find_matching_json(doc):
    cn = canonical_name(doc["FileName"])
    cd = doc["Directory"].split("/")[-1]
    jsonmatch = f".*{cd}/{cn}.*json"
    return db.flatjson.find_one(
        {
            "$and": [
                {"SourceFile": {"$regex": jsonmatch}},
            ]
        }
    )

def fix_missing_date_videos():
    BAD_DIR = "/Users/frahof/frank_import/tttt"
    extensions = [e.strip(".") for e in video_formats]
    missing_date = db.flatjson.find(
        {
            "$and": [
                # {"SourceFile": {"$regex": ".*takeout-20201108T103354Z-001.*"}},
                {"FileTypeExtension": {"$in": extensions}},
                {
                    "$or": [
                        # {"DateTimeOriginal": None},
                        {"CreateDate": None},
                        {"CompressorName": "Apple Intermediate Codec"},
                    ]
                },
            ]
        }
    )
    for m in missing_date:
        j = find_matching_json(m)
        if not j:
            print(f"moving bad file {m['SourceFile']}")
            shutil.move(m["SourceFile"], os.path.join(BAD_DIR, m["FileName"]))
            db.media.delete_one({"_id": ObjectId(m["_id"])})
            continue
        base, ext = os.path.splitext(m["SourceFile"])
        new_ext = ".mp4"
        print(
            f"ffmpeg -hide_banner -loglevel panic -i '{m['SourceFile']}' -c:v libx264 -c:a aac -pix_fmt yuv420p -y '{base}{new_ext}'"
        )
        print(
            f"exiftool -m -overwrite_original -CreationDate='{j['PhotoTakenTimeFormatted']}' -CreateDate='{j['PhotoTakenTimeFormatted']}' '{base}{new_ext}'"
        )
        print(f"rm '{m['SourceFile']}'")
        db.media.update_one(
            {"_id": m["_id"]},
            {
                "$set": {
                    "CreationDate": j["PhotoTakenTimeFormatted"],
                    "CreateDate": j["PhotoTakenTimeFormatted"],
                    "SourceFile": m["SourceFile"].replace(ext, new_ext),
                    "FileName": m["FileName"].replace(ext, new_ext),
                    "MIMEType": "video/mp4",
                }
            },
            upsert=False,
        )

def fix_missing_date_images():
    BAD_DIR = "/Users/frahof/frank_import/tttt"
    extensions = [e.strip(".") for e in image_formats]
    missing_date = db.flatjson.find(
        {
            "$and": [
                # {"SourceFile": {"$regex": ".*takeout-20201108T103354Z-002.*"}},
                {"FileTypeExtension": {"$in": extensions}},
                {"DateTimeOriginal": None},
                {"CreateDate": None},
            ]
        }
    )
    for m in missing_date:
        j = find_matching_json(m)
        if not j:
            print(f"moving bad file {m['SourceFile']}")
            shutil.move(m["SourceFile"], os.path.join(BAD_DIR, m["FileName"]))
            db.media.delete_one({"_id": ObjectId(m["_id"])})
            continue
        print(
            f"updating DateTimeOriginal of {m['FileName']} to {j['PhotoTakenTimeFormatted']}"
        )
        update_exif_date(m, j["PhotoTakenTimeFormatted"])
        db.media.update_one(
            {"_id": m["_id"]},
            {
                "$set": {
                    "DateTimeOriginal": j["PhotoTakenTimeFormatted"],
                    "CreateDate": j["PhotoTakenTimeFormatted"],
                }
            },
            upsert=False,
        )

if __name__ == "__main__":
    target_dir = "/Volumes/Photos/frank/"
    os.chdir(target_dir)
    client = MongoClient("mongodb://mongodb:27017/")
    db = client.iphoto
    media_collection = db.media

    # move_edited_to_original()
    # fix_wrong_heic()
    # fix_missing_date_images()
    fix_missing_date_videos()

    # for m in media_collection.find(
    #         {"MIMEType": {"$eq": None}}
    # ):
    #     mt = mimetypes.MimeTypes().guess_type(m["SourceFile"])[0]
    #     print(mt)
    #     db.media.update_one(
    #         {"_id": m["_id"]},
    #         {
    #             "$set": {
    #                 "MIMEType": mt
    #             }
    #         },
    #         upsert=False,
    #     )

    # move_subfolders_into_batches("/Volumes/Photos/frank/files")
    # doc = media_collection.find_one(
    #     {"SourceFile": "./takeout-20201108T103354Z-001/Takeout/Google Photos/2020-07-24 #2/IMG_0058.PNG"}
    # )
    # update_exif_date(doc, '2020-11-07 09:03:18')
    # remove duplicate images/videos


from datetime import datetime, timedelta

import evernote.edam.type.ttypes as EnTypes
from evernote.api.client import EvernoteClient

notebook_dct = {'!nbox': 'cd9cb2e4-452e-4900-b494-dba98dfcbc1e',
                'Cabinet': '2ce80038-4c1f-4588-8a8b-2e93034521fb',
                'Connections': 'a6f78bbf-bbd8-43dc-ad36-7501d2c083f9',
                'J-Tour': 'e19eb69d-6a49-47b8-b516-127e48c801ac',
                'Memories': '8f5e65d6-0dd6-443c-bf35-e0be151a1711',
                'Personal': '923944ba-3d61-4db9-86c7-17bf060e79a7'}

def convert_time(local_dt):
    delta = (local_dt - datetime.utcfromtimestamp(0)).total_seconds()
    # dirty hack to solve the local time to utc problem... work for pst
    delta += + 3600*12
    return int(delta * 1000)

class MyBook():
    def __init__(self, dev_tok_file='dev_token.txt'):
        devtok = open(dev_tok_file, 'r').read().strip()
        self.token = devtok
        self.client = EvernoteClient(token=devtok, sandbox=False)
        self.user_store = self.client.get_user_store()
        self.note_store = self.client.get_note_store()
        pass

    def create_note(self, title, content, created=None, notebook_name=None):
        note = EnTypes.Note()
        note.title = title
        note.content = content
        if created:
            note.created = convert_time(created)
        if notebook_name:
            note.notebookGuid = notebook_dct[notebook_name]
        created_note_guid = self.note_store.createNote(note)
        print('Created Note: %s' % created_note_guid)


# %% Importing libraries

import torch
import torch.nn as nn
import torch.optim as optim

import matplotlib.pyplot as plt
import numpy as np
import scipy.io as sio

import time
import os
import random
import copy

import loss_landscapes
import loss_landscapes.metrics

from mpl_toolkits.mplot3d import axes3d, Axes3D # ignore warning, we need this
from sklearn.decomposition import PCA

from CNN_functions import files_tensor, getListOfFiles, batch_generator, string_tensor, \
        sec_max, weights_init, plot_confusion_matrix, weights_to_list, list_to_weights

import pandas as pd
import seaborn as sb

from scipy.special import softmax

from evt_fitting_V2 import weibull_tailfitting
from compute_openmax_V2 import recalibrate_scores, computeOpenMaxProbability

# %%  ----- Delete old net -----------------------------------------------------------------------------------------------

if 'net' in locals():
    del net

# %%  ----- Set Parameters and Folders -----------------------------------------------------------------------------------

CUDA_flag = False

# Place where output folders are created:
DT = r'C:\Users\ddavidse\Desktop'

# Location of original data:
dirName = r'C:\SharedData\Datasets for CNN\100x100 downsampled from 150x150'

# Location of fooling data:
dataset_fool_dir = r'C:\SharedData\Datasets for CNN\elephant set'
#dataset_fool_dir = r'C:\SharedData\Datasets for CNN\mirrored set'

BS = 50             # batch size
epochs = 12         # number of epochs for training
fig_dpi = 150       # dpi for saved figures
fig_font_size = 12  # font size for figures

# %% ----- Control randomness -----------------------------------------------------------------------------------

random_initialization = True        # set False to use saved_weights
saved_weights = r'D:\ddavidse\Desktop\5 class network runs\network runs - background - full data\Network_run_16\initial_weights.pth'

set_random_seed = False             # set True to use seed_value
seed_value = 784410

# %% ----- Configure dataset split ---------------------------------------------------------------------------------------

ValFrac = 1/6       # set fraction of total dataset used for validation
TestFrac = 1/6      # set fraction of total dataset used for testing

# %% ----- Select features -----------------------------------------------------------------------------------------------

batchnorm_flag = True                   # make True to use batch normalization
misclassified_images_flag = False       # make True to output misclassified images
misclassified_outputs_flag = False      # make True to output misclassified net output values to run info.txt
loss_landscape_flag = False             # make True to create a loss landscape based on random planes
PCA_flag = False                        # make True to create a loss landscape based on PCA

baseline_flag = False                   # make True to use probability thresholding
MAV_flag = True                         # make True to do all the MAV magic
openmax_flag = True                     # make True to use OpenMax (requires MAV_flag to be true as well)

# %% ----- Loss landscapes package settings ------------------------------------------------------------------------------

STEPS = 40      # resolution of landscape, higher = more detail

# %% ----- PCA settings --------------------------------------------------------------------------------------------------

FilterSteps = 40
filter_normalization_flag = True
distance_multiplier = 2
bowl = False
contour_center = False

# %%  ----- CUDA initialization ------------------------------------------------------------------------------------------

if CUDA_flag:
    torch.set_default_tensor_type('torch.cuda.FloatTensor')

# %% ----- Create output folder ------------------------------------------------------------------------------------------

os.chdir(DT)
mfname = 'Network_runs'

if not os.path.exists(mfname):
    os.mkdir(mfname)

fdir = '{}\\{}'.format(DT,mfname)

counter = 0
tv = 0

while tv == 0:

    counter += 1
    fname = '{}\\Network_run_{}'.format(fdir, counter)

    if not os.path.exists(fname):
        os.mkdir(fname)
        tv = 1

# %% ----- detecting number of images per class --------------------------------------------------------------------------

os.chdir(dirName)

DirListTop = os.listdir()
names = DirListTop

N = []

for x in names:

    os.chdir(x)
    DirListSub = os.listdir()
    N.append(len(DirListSub))
    os.chdir(dirName)

NC = np.cumsum(N)

# %% ----- detecting file size -------------------------------------------------------------------------------------------

os.chdir(names[0])
filelist = os.listdir()
testfile = filelist[0]

testimage = sio.loadmat(testfile)
testimage2 = list(testimage.values())
testimage3 = testimage2[3]

input_size = len(testimage3)

# %% ----- Obtaining and storing data ------------------------------------------------------------------------------------

# Data
listOfFiles = getListOfFiles(dirName)
listOfImages = files_tensor(listOfFiles)

# Labels
label = np.zeros(sum(N))
labels = [x for x in range(len(names))]

for i in range(0,NC[0]):
    label[i] = labels[0]

for k in range(1,len(N)):
    for i in range(NC[k-1],NC[k]):
        label[i] = labels[k]

# Data and labels
shuffle_list = list(zip(listOfImages, label))
dataset_orig = [listOfImages, label]

# %% ----- Creating sets with equal class distribution -------------------------------------------------------------------

if set_random_seed == True:
    SeedVal = seed_value
else:
    SeedVal = random.randrange(1,1000000)

random.seed(SeedVal)

zipZ = [[] for x in range(len(N))]
zip_train = [[] for x in range(len(N))]
zip_val = [[] for x in range(len(N))]
zip_test = [[] for x in range(len(N))]

zipZ[0] = shuffle_list[0:NC[0]]

for k in range(1,len(N)):
    zipZ[k] = shuffle_list[NC[k-1]:NC[k]]

for k in range(len(N)):
    random.shuffle(zipZ[k])

TrainFrac = 1 - ValFrac - TestFrac

for k in range(len(N)):
    zip_train[k] = zipZ[k][0:int(TrainFrac*N[k])]
    zip_val[k] = zipZ[k][int(TrainFrac*N[k]):int((TrainFrac+ValFrac)*N[k])]
    zip_test[k] = zipZ[k][int((TrainFrac+ValFrac)*N[k]):N[k]]

train = [x for s in zip_train for x in s]
val = [x for s in zip_val for x in s]
test = [x for s in zip_test for x in s]

files1, labels1 = zip(*test)
test_set = [files1, labels1]

files2, labels2 = zip(*train)
train_set = [files2, labels2]

files3, labels3 = zip(*val)
val_set = [files3, labels3]

# %% ----- Adjusting sets to fit in batches (optional code, not important) -----------------------------------------------

diff_val = BS - (len(val_set[0]) % BS);

if diff_val == 2:

    c_images = []
    c_labels = []

    c_images.append(train_set[0][0])
    c_labels.append(train_set[1][0])

    c_images.append(train_set[0][len(train_set[0]) - 1])
    c_labels.append(train_set[1][len(train_set[1]) - 1])

    c_images = tuple(c_images)
    c_labels = tuple(c_labels)

    val_set[0] = val_set[0] + c_images
    val_set[1] = val_set[1] + c_labels

    train_set[0] = train_set[0][1:len(train_set[0])-1]
    train_set[1] = train_set[1][1:len(train_set[1])-1]

elif diff_val == 1:

    c_image = train_set[0][len(train_set[0]) - 1]
    c_label = train_set[1][len(train_set[1]) - 1]

    t0list = list(val_set[0])
    t0list.append(c_image)
    val_set[0] = tuple(t0list)

    t1list = list(val_set[1])
    t1list.append(c_label)
    val_set[1] = tuple(t1list)

    train_set[0] = train_set[0][0:len(train_set[0])-2]
    train_set[1] = train_set[1][0:len(train_set[1])-2]

diff_test = BS - (len(test_set[0]) % BS);

if diff_test == 2:

    c_images = []
    c_labels = []

    c_images.append(train_set[0][0])
    c_labels.append(train_set[1][0])

    c_images.append(train_set[0][len(train_set[0]) - 1])
    c_labels.append(train_set[1][len(train_set[1]) - 1])

    c_images = tuple(c_images)
    c_labels = tuple(c_labels)

    test_set[0] = test_set[0] + c_images
    test_set[1] = test_set[1] + c_labels

    train_set[0] = train_set[0][1:len(train_set[0])-1]
    train_set[1] = train_set[1][1:len(train_set[1])-1]

elif diff_test == 1:

    c_image = train_set[0][len(train_set[0]) - 1]
    c_label = train_set[1][len(train_set[1]) - 1]

    t0list = list(test_set[0])
    t0list.append(c_image)
    test_set[0] = tuple(t0list)

    t1list = list(test_set[1])
    t1list.append(c_label)
    test_set[1] = tuple(t1list)

    train_set[0] = train_set[0][0:len(train_set[0])-1]
    train_set[1] = train_set[1][0:len(train_set[1])-1]

diff_train = len(test_set[0]) % BS

if diff_train < 5:

    c_image = train_set[0][len(test_set[0]) - 1]
    c_label = train_set[1][len(test_set[1]) - 1]

    t0list = list(train_set[0])
    t0list.append(c_image)
    train_set[0] = tuple(t0list)

    t1list = list(train_set[1])
    t1list.append(c_label)
    train_set[1] = tuple(t1list)

    test_set[0] = test_set[0][0:len(test_set[0])-1]
    test_set[1] = test_set[1][0:len(test_set[1])-1]

# %% ----- Batch generator and loaders ----------------------------------------------------------------------------------

train_loader = batch_generator (BS, train_set)
test_loader = batch_generator (BS, test_set)
val_loader = batch_generator (BS, val_set)

# %% ----- Soft coding input size ----------------------------------------------------------------------------------------

if input_size % 2 == 0:
    Size_1 = input_size / 2
else:
    Size_1 = (input_size - 1) / 2

if Size_1 % 2 == 0:
    Size_2 = Size_1 / 2
else:
    Size_2 = (Size_1 - 1) / 2

Size_2 = int(Size_2)

# %% ----- Defining the neural nets --------------------------------------------------------------------------------------

#note: nn.Conv2D(in, out, kernel, stride=1, padding)
#note: nn.MaxPool2d(kernel, stride, padding)

class BN_Net(nn.Module):

    def __init__(self):
        super(BN_Net, self).__init__()

        self.features = nn.Sequential(
                                      nn.Conv2d(1,5,5,1,2),
                                      nn.MaxPool2d(2,2),
                                      nn.ReLU(inplace=True),
                                      nn.BatchNorm2d(5),
                                      nn.Conv2d(5,8,5,1,2),
                                      nn.MaxPool2d(2,2),
                                      nn.ReLU(inplace=True),
                                      nn.BatchNorm2d(8)
                                         )

        self.classifier = nn.Sequential(
                                        nn.Linear(8*Size_2*Size_2, 120),
                                        nn.ReLU(inplace=True),
                                        nn.BatchNorm1d(120),
                                        nn.Linear(120,84),
                                        nn.ReLU(inplace=True),
                                        nn.BatchNorm1d(84),
                                        nn.Linear(84,5)
                                        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(-1, 8*Size_2*Size_2)
        x = self.classifier(x)

        return x

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        self.features = nn.Sequential(
                                      nn.Conv2d(1,5,5,1,2),
                                      nn.MaxPool2d(2,2),
                                      nn.ReLU(inplace=True),
                                      nn.Conv2d(5,8,5,1,2),
                                      nn.MaxPool2d(2,2),
                                      nn.ReLU(inplace=True)
                                         )

        self.classifier = nn.Sequential(
                                        nn.Linear(8*Size_2*Size_2, 120),
                                        nn.ReLU(inplace=True),
                                        nn.Linear(120,84),
                                        nn.ReLU(inplace=True),
                                        nn.Linear(84,5)
                                        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(-1, 8*Size_2*Size_2)
        x = self.classifier(x)

        return x

# %% ----- Initialize or load weights -----------------------------------------------------------------------------

if batchnorm_flag:
    net = BN_Net()
else:
    net = Net()

if CUDA_flag:
    net.cuda()

if random_initialization:
    net.apply(weights_init)
else:
    checkpoint = torch.load(saved_weights)
    net.load_state_dict(checkpoint['model_state_dict'])

# %% ----- Loss and optimizer ---------------------------------------------------------------------------------------

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(),lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False)
#optimizer = optim.RMSprop(net.parameters(), lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)
#optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0, dampening=0, weight_decay=0, nesterov=False)
model_initial = copy.deepcopy(net)

# %% ----- Training the net ---------------------------------------------------------------------------------------------

start = time.time()

running_loss_history = []
running_corrects_history = []
val_running_loss_history = []
val_running_corrects_history = []

StateVecList = []

for e in range(epochs):

    print('\nepoch :', (e+1))
    running_loss = 0.0
    running_corrects = 0.0
    val_running_loss = 0.0
    val_running_corrects = 0.0
    train_loader = batch_generator(BS, train_set)
    val_loader = batch_generator(BS, val_set)

    train_amount = 0.0
    val_amount = 0.0

    for inputs, labels in train_loader:

        labels = string_tensor(labels)
        inputs = torch.stack(inputs)
        inputs = torch.unsqueeze(inputs, 1)
        outputs = net(inputs)
        loss = criterion(outputs, labels.long())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        _, preds = torch.max(outputs, 1)
        running_loss += loss.item()
        running_corrects += torch.sum(preds == labels.long().data)

        train_amount += BS

    if PCA_flag:
        #StateVecList.append(weights_to_list(net))
        print('nvm')

    with torch.no_grad():
        for val_inputs, val_labels in val_loader:

            val_inputs = torch.stack(val_inputs)
            val_inputs = torch.unsqueeze(val_inputs, 1)
            val_labels = string_tensor(val_labels)
            val_outputs = net(val_inputs)
            val_loss = criterion(val_outputs, val_labels.long())

            _, val_preds = torch.max(val_outputs, 1)
            val_running_loss += val_loss.item()
            val_running_corrects += torch.sum(val_preds == val_labels.long().data)

            val_amount += BS

    epoch_loss = running_loss
    epoch_acc = running_corrects
    running_loss_history.append(epoch_loss)
    running_corrects_history.append(epoch_acc)
    val_epoch_loss = val_running_loss
    val_epoch_acc = val_running_corrects
    val_running_loss_history.append(val_epoch_loss)
    val_running_corrects_history.append(val_epoch_acc)

    print('\ntraining loss: {:.4f} '.format(epoch_loss))
    print('validation loss: {:.4f}'.format(val_epoch_loss))

if CUDA_flag:
    torch.cuda.synchronize()

end = time.time()
print('\n-----------------------------------------------------------------------')
print('\nTime of training: {:d} s'.format(round((end - start))))

# %% ----- Prediction and training stats and graphs --------------------------------------------------------------------

plt.rcParams.update({'font.size': fig_font_size})

correct = [0,0,0]
total = [0,0,0]
testsize = [0,0,0]
AV_all = [[],[],[],[],[]]

with torch.no_grad():
    for data in test_loader:
        images_0_0, labels_0 = data
        images_0_1 = torch.stack(images_0_0)
        images_0_2 = torch.unsqueeze(images_0_1, 1)
        outputs_0 = net(images_0_2)
        _, predicted_0 = torch.max(outputs_0.data, 1)
        total[0] += string_tensor(labels_0).size(0)
        testsize[0] += 1

        correct_index_0 = predicted_0 == string_tensor(labels_0).long()
        correct[0] += correct_index_0.sum().item()

        if MAV_flag:
            for x in range(len(outputs_0)):
                if correct_index_0[x]:
                    AV_all[predicted_0[x]].append(outputs_0[x])

    train_loader = batch_generator(BS, train_set)
    val_loader = batch_generator(BS, val_set)

    for data in train_loader:
        images_1_0, labels_1 = data
        images_1_1 = torch.stack(images_1_0)
        images_1_2 = torch.unsqueeze(images_1_1, 1)
        outputs_1 = net(images_1_2)
        _, predicted_1 = torch.max(outputs_1.data, 1)
        total[1] += string_tensor(labels_1).size(0)
        testsize[1] += 1

        correct_index_1 = predicted_1 == string_tensor(labels_1).long()
        correct[1] += correct_index_1.sum().item()

        if MAV_flag:
            for x in range(len(outputs_1)):
                if correct_index_1[x]:
                    AV_all[predicted_1[x]].append(outputs_1[x])

    for data in val_loader:
        images_2_0, labels_2 = data
        images_2_1 = torch.stack(images_2_0)
        images_2_2 = torch.unsqueeze(images_2_1, 1)
        outputs_2 = net(images_2_2)
        _, predicted_2 = torch.max(outputs_2.data, 1)
        total[2] += string_tensor(labels_2).size(0)
        testsize[2] += 1

        correct_index_2 = predicted_2 == string_tensor(labels_2).long()
        correct[2] += correct_index_2.sum().item()

        if MAV_flag:
            for x in range(len(outputs_2)):
                if correct_index_2[x]:
                    AV_all[predicted_2[x]].append(outputs_2[x])

if MAV_flag:

    MAV = np.empty((5,5))

    for i in range(5):
        AV_array = np.array([list(x.cpu().numpy()) for x in AV_all[i]])
        MAV[i] = np.mean(AV_array, axis=0)

"""
================================================
================================================
===== SECONDARY DATASET LOOPS FOR MAV
================================================
================================================
"""

AV_dist_correct = [[],[],[],[],[]]
AV_dist_wrong = []

if MAV_flag:

    train_loader = batch_generator(BS, train_set)
    val_loader = batch_generator(BS, val_set)
    test_loader = batch_generator(BS, test_set)

    with torch.no_grad():
        for data in test_loader:
            images_0_0, labels_0 = data
            images_0_1 = torch.stack(images_0_0)
            images_0_2 = torch.unsqueeze(images_0_1, 1)
            outputs_0 = net(images_0_2)
            _, predicted_0 = torch.max(outputs_0.data, 1)

            correct_index_0 = predicted_0 == string_tensor(labels_0).long()

            for x in range(len(outputs_0)):
                avdiff = outputs_0[x].cpu().numpy() - MAV[predicted_0[x]]
                avdist = np.linalg.norm(avdiff)

                if correct_index_0[x]:
                    AV_dist_correct[predicted_0[x]].append(avdist)
                else:
                    AV_dist_wrong.append(avdist)

        for data in train_loader:
            images_1_0, labels_1 = data
            images_1_1 = torch.stack(images_1_0)
            images_1_2 = torch.unsqueeze(images_1_1, 1)
            outputs_1 = net(images_1_2)
            _, predicted_1 = torch.max(outputs_1.data, 1)

            correct_index_1 = predicted_1 == string_tensor(labels_1).long()

            for x in range(len(outputs_1)):
                avdiff = outputs_1[x].cpu().numpy() - MAV[predicted_1[x]]
                avdist = np.linalg.norm(avdiff)

                if correct_index_1[x]:
                    AV_dist_correct[predicted_1[x]].append(avdist)
                else:
                    AV_dist_wrong.append(avdist)

        for data in val_loader:
            images_2_0, labels_2 = data
            images_2_1 = torch.stack(images_2_0)
            images_2_2 = torch.unsqueeze(images_2_1, 1)
            outputs_2 = net(images_2_2)
            _, predicted_2 = torch.max(outputs_2.data, 1)

            correct_index_2 = predicted_2 == string_tensor(labels_2).long()

            for x in range(len(outputs_0)):
                avdiff = outputs_2[x].cpu().numpy() - MAV[predicted_2[x]]
                avdist = np.linalg.norm(avdiff)

                if correct_index_2[x]:
                    AV_dist_correct[predicted_2[x]].append(avdist)
                else:
                    AV_dist_wrong.append(avdist)

    AV_dist_tot = [x for i in AV_dist_correct for x in i]

    AVCA = np.array(AV_dist_tot)
    AVC_max = np.max(AVCA)
    AVC_min = np.min(AVCA)
    AVC_avg = np.mean(AVCA)
    AVC_std = np.std(AVCA)

    AVWA = np.array(AV_dist_wrong)
    AVW_max = np.max(AVWA)
    AVW_min = np.min(AVWA)
    AVW_avg = np.mean(AVWA)
    AVW_std = np.std(AVWA)

    test_loader = batch_generator(BS, test_set)

    with torch.no_grad():

        ImagesChecked = 0
        ImagesCheckedCorrect = 0

        for data in test_loader:
            images_3_0, labels_3 = data
            images_3_1 = torch.stack(images_3_0)
            images_3_2 = torch.unsqueeze(images_3_1, 1)
            outputs_3 = net(images_3_2)
            _, predicted_3 = torch.max(outputs_3.data, 1)

            correct_index_3 = predicted_3 == string_tensor(labels_3).long()

            for x in range(len(outputs_3)):
                avdiff = outputs_3[x].cpu().numpy() - MAV[predicted_3[x]]
                avdist = np.linalg.norm(avdiff)
                if avdist < AVW_min:
                #if avdist < 2.5:
                    ImagesChecked += 1
                    if correct_index_3[x]:
                        ImagesCheckedCorrect += 1

    df1 = pd.DataFrame({'1':AVCA})
    df2 = pd.DataFrame({'2':AVWA})
    df = pd.concat([df1, df2], ignore_index=True, axis=1)
    df.columns = ['correct','incorrect']

    MAV_dist_plot = plt.figure(figsize=(9,7))
    plt.grid(1, which='major')
    plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
    plt.minorticks_on()
    sbplot = sb.stripplot(data=df)
    plt.ylabel('distance to MAV')
    x = plt.gca().axes.get_xlim()
    plt.plot(x, len(x)*[AVW_min],'r')

# %%

"""
================================================
================================================
===== FOOLING DATASET
================================================
================================================

"""

if MAV_flag:

    thresh_flag = True
    thresh = 4.7

    os.chdir(dataset_fool_dir)

    DirListTop_fool = os.listdir()
    names_fool = DirListTop_fool

    N_fool = []

    for x in names_fool:

        os.chdir(x)
        DirListSub_fool = os.listdir()
        N_fool.append(len(DirListSub_fool))
        os.chdir(dataset_fool_dir)

    NC_fool = np.cumsum(N_fool)

    listOfFiles_fool = getListOfFiles(dataset_fool_dir)
    listOfImages_fool = files_tensor(listOfFiles_fool)

    label_fool = np.zeros(sum(N_fool))
    labels_fool = [x for x in range(len(names_fool))]

    for i in range(0,NC_fool[0]):
        label_fool[i] = labels_fool[0]

    for k in range(1,len(N_fool)):
        for i in range(NC_fool[k-1],NC_fool[k]):
            label_fool[i] = labels_fool[k]

    dataset_fool = [listOfImages_fool, label_fool]

    list_1 = AV_dist_tot
    list_2 = AV_dist_wrong

    data_loader = batch_generator(BS, dataset_fool)

    dist_test = []

    with torch.no_grad():
        for data in data_loader:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in range(len(outputs)):
                avdiff = outputs[x].cpu().numpy() - MAV[predicted[x]]
                avdist = np.linalg.norm(avdiff)

                dist_test.append(avdist)

    df0 = pd.DataFrame({'0':list_1 + list_2})
    df1 = pd.DataFrame({'1':list_1})
    df2 = pd.DataFrame({'2':list_2})
    df3 = pd.DataFrame({'3':dist_test})

    df = pd.concat([df0, df3], ignore_index=True, axis=1)
    df.columns = ['original set','fooling set']

    MAV_dist_plot = plt.figure(figsize=(9,7))
    plt.grid(1, which='major')
    plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
    plt.minorticks_on()
    sbplot = sb.stripplot(data=df, jitter=0.15)
    plt.ylabel('distance to MAV')

    if thresh_flag:

        x = plt.gca().axes.get_xlim()
        plt.plot(x, len(x)*[thresh],'r')

        list_1_thr = [x for x in list_1 if x < thresh]
        list_2_thr = [x for x in list_2 if x < thresh]
        dist_test_thr = [x for x in dist_test if x < thresh]

        L1C = len(list_1)
        L1W = len(list_2)
        L1T = len(dist_test)

        L2C = len(list_1_thr)
        L2W = len(list_2_thr)
        L2T = len(dist_test_thr)

        A1 = round(100 * L1C / (L1C + L1W), 1)
        A2 = round(100 * L2C / (L2C + L2W), 1)

        P1 = round(100 * (1 - (L2C + L2W) / (L1C + L1W)), 1)
        P2 = round(100 * (1 - L2T / L1T), 1)

        print('\n-------------------------------------------------------------------')

        os.chdir(fname)

        path = '{}\\MAV_distance_dist.png'.format(fname)
        MAV_dist_plot.savefig(path, dpi=fig_dpi)

        f = open('MAV results.txt','w+')
        f.write('--------------------------------------------------------------------------------------')
        f.write('\n\nAccuracy on the original dataset: \t\t{} %'.format(A1))
        f.write('\nAccuracy on thresholded original dataset: \t{} %'.format(A2))
        f.write('\n\nThreshold: {}'.format(thresh))
        f.write('\n\nPart of the original dataset thrown away: \t{} %'.format(P1))
        f.write('\nPart of the secondary dataset thrown away: \t{} %'.format(P2))
        f.write('\n\n--------------------------------------------------------------------------------------')
        f.close()

# %%

"""
================================================
================================================
===== OPENMAX
================================================
================================================
"""

if openmax_flag:

    tail = 2
    open_thresh = 0.2
    alpha = 4

    maxarray = np.array([1,2,3,4,5])

    weibull_models = weibull_tailfitting(AV_dist_correct, MAV, names, tailsize=tail, distance_type='euclidean')

    data_loader_fool = batch_generator(BS, dataset_fool)

    n_open_set = 0
    n_set = 0
    SMM_fool = []
    SMM_orig = []

    with torch.no_grad():
        for data in data_loader_fool:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in outputs:
                openmax_probs = recalibrate_scores(weibull_models, x.numpy(), alpharank=alpha)
                n_set += 1

                SMM = max(openmax_probs[:-1])
                SMM_fool.append(SMM)

                if SMM < open_thresh:
                    n_open_set += 1

    open_throw = round(100 * n_open_set / n_set, 1)

    data_loader_orig = batch_generator(BS, dataset_orig)

    with torch.no_grad():
        for data in data_loader_orig:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in outputs:
                openmax_probs = recalibrate_scores(weibull_models, x.numpy(), alpharank=alpha)
                n_set += 1

                SMM = max(openmax_probs[:-1])
                SMM_orig.append(SMM)

                if SMM < open_thresh:
                    n_open_set += 1

    open_throw_2 = round(100 * n_open_set / n_set, 1)

    df1 = pd.DataFrame({'1':SMM_orig})
    df2 = pd.DataFrame({'2':SMM_fool})
    df = pd.concat([df1, df2], ignore_index=True, axis=1)
    df.columns = ['original','fool']

    openmax_plot = plt.figure(figsize=(9,7))
    plt.grid(1, which='major')
    plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
    plt.minorticks_on()
    sbplot = sb.stripplot(data=df)
    plt.ylabel('max prob 1:5')
    x = plt.gca().axes.get_xlim()
    plt.plot(x, len(x)*[open_thresh],'r')

    print('\n-------------------------------------------------------------------------------------')
    print('\nOpenMax threw out {} % of the original dataset'.format(open_throw_2))
    print('OpenMax threw out {} % of the fooling dataset'.format(open_throw))
    print('\nDifference: {:.1f} %'.format(open_throw - open_throw_2))

    os.chdir(fname)

    path = '{}\openmax_plot.png'.format(fname)
    openmax_plot.savefig(path, dpi=fig_dpi)

    f = open('openmax info.txt','w+')
    f.write('-----------------------------------------------------------------------------------------')
    f.write('\n\nparameters used:')
    f.write('\n\ntail = {}'.format(tail))
    f.write('\nopen_thresh = {}'.format(open_thresh))
    f.write('\nalpha = {}'.format(alpha))
    f.write('\n\n-----------------------------------------------------------------------------------------')
    f.write('\n\nResults:')
    f.write('\n\nOpenMax threw out {} % of the fooling dataset'.format(open_throw))
    f.write('\nOpenMax threw out {} % of the original dataset'.format(open_throw_2))
    f.write('\n\nDifference: {:.1f} %'.format(open_throw - open_throw_2))
    f.write('\n\n-----------------------------------------------------------------------------------------')
    f.close()

    os.chdir(DT)

# %%

"""
================================================
================================================
===== SECONDARY DATASET LOOPS FOR BASELINE
================================================
================================================
"""

if baseline_flag:

    train_loader = batch_generator(BS, train_set)
    val_loader = batch_generator(BS, val_set)
    test_loader = batch_generator(BS, test_set)

    unc_correct = []
    unc_wrong = []

    with torch.no_grad():

        for data in test_loader:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in range(len(outputs)):

                probs = softmax(outputs[x].cpu().numpy())
                prob = np.max(probs)
                uncert = 100*(1 - prob)

                if correct_index[x]:
                    unc_correct.append(uncert)
                else:
                    unc_wrong.append(uncert)

        for data in train_loader:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in range(len(outputs)):

                probs = softmax(outputs[x].cpu().numpy())
                prob = np.max(probs)
                uncert = 100*(1 - prob)

                if correct_index[x]:
                    unc_correct.append(uncert)
                else:
                    unc_wrong.append(uncert)

        for data in val_loader:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in range(len(outputs)):

                probs = softmax(outputs[x].cpu().numpy())
                prob = np.max(probs)
                uncert = 100*(1 - prob)

                if correct_index[x]:
                    unc_correct.append(uncert)
                else:
                    unc_wrong.append(uncert)

    df1 = pd.DataFrame({'1':unc_correct})
    df2 = pd.DataFrame({'2':unc_wrong})
    df = pd.concat([df1, df2], ignore_index=True, axis=1)
    df.columns = ['correct','wrong']

    baseline_plot = plt.figure(figsize=(9,7))
    plt.grid(1, which='major')
    plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
    plt.minorticks_on()
    sbplot = sb.stripplot(data=df, jitter=.15)
    plt.ylabel('uncertainty in %')
    plt.title('baseline uncertainty plot')

total_tot = sum(total)
correct_tot = sum(correct)
testsize_tot = sum(testsize)

print('\nAccuracy of the network on the test images: {:.1f} %'.format((100 * correct[0] / total[0])))
print('Test result obtained from {} images coming from {} batches'.format(total[0], testsize[0]))
print('\nAccuracy of the network on the total set of all images: {:.1f} %'.format((100 * correct_tot / total_tot)))
print('Test result obtained from {} images coming from {} batches'.format(total_tot, testsize_tot))

if MAV_flag:
    print('\nAmount of images in high confidence set: {}'.format(ImagesChecked))
    print('This is {:.1f}% of the test set'.format(100*ImagesChecked / total[0]))

    print('\nAmount of correctly classified images in high confidence set: {}'.format(ImagesCheckedCorrect))
    print('Accuracy on high confidence set: {:.1f} %'.format(100*ImagesCheckedCorrect / ImagesChecked))

#print('\nAmount of correct predictions in test 1: {}'.format(correct[0]))

if MAV_flag:
    print('\n-------------------------------------------------------------------------------------')
    print('\nAccuracy on the original dataset: \t\t{} %'.format(A1))
    print('Accuracy on thresholded original dataset: \t{} %'.format(A2))
    print('\nThreshold: {}'.format(thresh))
    print('\nPart of the original dataset thrown away: \t{} %'.format(P1))
    print('Part of the secondary dataset thrown away: \t{} %'.format(P2))
    if openmax_flag:
        print('\n-------------------------------------------------------------------------------------')
        print('\nOpenMax threw out {} % of the original dataset'.format(open_throw_2))
        print('OpenMax threw out {} % of the fooling dataset'.format(open_throw))
        print('\nDifference: {:.1f} %'.format(open_throw - open_throw_2))

val_plot = plt.figure(figsize=(9,7))
plt.grid(1, which='major')
plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
plt.minorticks_on()
plt.plot(running_loss_history, 'r',label='training loss')
plt.plot(val_running_loss_history,'b', label='validation loss')
plt.autoscale(enable=True, axis='x', tight=True)
plt.legend()

v_len = len(val_set[0])
t_len = len(train_set[0])

acc_plot = plt.figure(figsize=(9,7))
plt.grid(1, which='major')
plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
plt.minorticks_on()
train_corrected = [100*float(x)/(train_amount) for x in running_corrects_history]
plt.plot(train_corrected,'r', label='training accuracy [%]')
val_corrected = [100*float(x)/(val_amount) for x in val_running_corrects_history]
plt.plot(val_corrected,'b', label='validation accuracy [%]')
plt.autoscale(enable=True, axis='x', tight=True)
plt.legend()

# %% ----- confusion matrix --------------------------------------------------------------------------------------------
# code inspiration from https://deeplizard.com/learn/video/0LhiS6yu2qQ

test_loader = batch_generator (BS, test_set)
all_preds = torch.tensor([])
all_labels = torch.tensor([])

misc_preds_t = []
misc_labels_t = []
misc_img_t = []
misc_outputs_t = []
maxdif_t = []

for inputs, labels in test_loader:

        inputs = torch.stack(inputs)
        inputs = torch.unsqueeze(inputs, 1)

        labels = string_tensor(labels)
        all_labels = torch.cat((all_labels, labels),dim=0)

        outputs = net(inputs)
        all_preds = torch.cat((all_preds, outputs),dim=0)

        imlist = list(inputs)
        outputs1 = outputs.argmax(dim=1)
        vec_equal = outputs1 == labels

        if misclassified_outputs_flag:

            misc_outputs_0 = outputs[~vec_equal]
            misc_outputs_1 = misc_outputs_0.cpu().detach().numpy().tolist()
            misc_outputs = [list(map(lambda x: round(x,1), i)) for i in misc_outputs_1]
            misc_outputs_t += misc_outputs

            maxdif = [max(x) - sec_max(x) for x in misc_outputs_1]
            maxdif_t += maxdif

        if misclassified_images_flag:

            misc_ind = [x for x in range(BS) if vec_equal[x] == False]

            misc_img = [imlist[x] for x in misc_ind]
            misc_img_t += misc_img

            misc_labels = [labels[x] for x in misc_ind]
            misc_labels_t += misc_labels

            misc_preds = [outputs1[x] for x in misc_ind]
            misc_preds_t += misc_preds

AL = all_labels
AL2 = np.array(list(AL), dtype=np.int)
AL3 = torch.tensor(AL2)

TP = all_preds.argmax(dim=1)
TP2 = np.array(list(TP), dtype=np.int)
TP3 = torch.tensor(TP2)

ConfAcc = sum(TP3 == AL3).item()
#print('Amount of correct predictions in test 2: {}'.format(ConfAcc))
print('\n-----------------------------------------------------------------------------------')

stacked = torch.stack((AL3, TP3), dim=1)

cmt = torch.zeros(len(N),len(N), dtype=torch.int64)

for p in stacked:
    tl, pl = p.tolist()
    cmt[tl, pl] = cmt[tl, pl] + 1

cmt2 = cmt.cpu()

if len(names) == 3:
    conf_fig_size = (8,8)
elif len(names) == 4:
    conf_fig_size = (9,9)
else:
    conf_fig_size = (10,10)

conf_fig = plt.figure(figsize = conf_fig_size)
plot_confusion_matrix(cmt2, names)

conf_fig_2 = plt.figure(figsize = conf_fig_size)
plot_confusion_matrix(cmt2, names, normalize='row', percentage=True)

conf_fig_3 = plt.figure(figsize = conf_fig_size)
plot_confusion_matrix(cmt2, names, normalize='full', percentage=True)

# %% ----- loss landscape ----------------------------------------------------------------------------------------------
# source: https://github.com/marcellodebernardi/loss-landscapes/blob/master/examples/core-features.ipynb

if loss_landscape_flag:
    start = time.time()

    train_loader = batch_generator (BS, train_set)
    x, y = iter(train_loader).__next__()
    x = torch.stack(x)
    x = torch.unsqueeze(x, 1)
    y = torch.tensor(y).long()
    metric = loss_landscapes.metrics.Loss(criterion, x, y)

    LCP = loss_landscapes.random_plane(net, metric, 100, STEPS, normalization='filter', deepcopy_model=True)

    if CUDA_flag:
        torch.cuda.synchronize()

    end = time.time()
    print('\nTime of calculating loss landscape: {:d} s'.format(round((end - start))))

    loss_con = plt.figure()
    plt.contour(LCP, levels=50)
    plt.title('Loss Contours around Trained Model')
    plt.show()

    loss_surf_1 = fig = plt.figure(figsize=(9,7))
    ax = plt.axes(projection='3d')

    X = np.array([[j for j in range(STEPS)] for i in range(STEPS)])
    Y = np.array([[i for _ in range(STEPS)] for i in range(STEPS)])
    ax.plot_surface(X, Y, LCP, rstride=1, cstride=1, cmap='coolwarm', edgecolor='none')
    ax.set_title('Surface Plot of Loss Landscape')
    ax.set_xlabel(r'$\theta$', fontsize=18, labelpad=10)
    ax.set_ylabel(r"$\theta '$", fontsize=18, labelpad=10)
    ax.set_zlabel('Loss', fontsize=18, labelpad=10, rotation=90)

    loss_surf_2 = fig = plt.figure(figsize=(9,7))
    ax = plt.axes(projection='3d')

    ax.plot_surface(X, Y, LCP, rstride=1, cstride=1, cmap='coolwarm', edgecolor='none')
    #ax.set_title('Surface Plot of Loss Landscape')
    ax.view_init(30, 45)
    ax.set_xlabel(r'$\theta$', fontsize=18, labelpad=10)
    ax.set_ylabel(r"$\theta '$", fontsize=18, labelpad=10)
    ax.set_zlabel('Loss', fontsize=18, labelpad=10, rotation=90)

# %% ----- PCA ----------------------------------------------------------------------------------------------------------

if PCA_flag:

    PCA_time_start = time.time()
    StateVecArray = np.array(StateVecList)
    TrainedNetVector = StateVecArray[epochs-1]

    pca = PCA(n_components=2)
    PC = pca.fit_transform(StateVecArray)
    PC_norm = []

    for i in range(len(PC[0])):
        col = PC[:,i]
        col_fixed = col / np.linalg.norm(col)
        PC_norm.append(col_fixed)

    SVA_1 = []
    for i in range(len(StateVecArray[0])):
        weightvec = StateVecArray[:,i]
        SVA_1.append(np.dot(weightvec, PC_norm[0]))

    SVA_2 = []
    for i in range(len(StateVecArray[0])):
        weightvec = StateVecArray[:,i]
        SVA_2.append(np.dot(weightvec, PC_norm[1]))

    SVA_1 = np.array(SVA_1)
    SVA_2 = np.array(SVA_2)

    if filter_normalization_flag:

        A = net.state_dict()
        S = [x.size() for x in A.values()]

        SL = [list(x) for x in S]

        L = len(SL)
        Sizes = []
        WeightTensors = []
        ite = 0

        for x in range(L):
            ite += 1
            xsize = 1
            for number in SL[x]:
                xsize = xsize * number
            Sizes.append(xsize)
            Sizes_C = np.cumsum(Sizes)

        Seg0 = TrainedNetVector[0:Sizes_C[0]]
        Seg0_Norm = np.linalg.norm(Seg0)

        Seg1 = SVA_1[0:Sizes_C[0]]
        Seg1_Norm = np.linalg.norm(Seg1)

        Seg2 = SVA_2[0:Sizes_C[0]]
        Seg2_Norm = np.linalg.norm(Seg2)

        SVA_1[0:Sizes_C[0]] = SVA_1[0:Sizes_C[0]] * Seg0_Norm / Seg1_Norm
        SVA_2[0:Sizes_C[0]] = SVA_2[0:Sizes_C[0]] * Seg0_Norm / Seg2_Norm

        TestNorm0 = np.linalg.norm(TrainedNetVector[0:Sizes_C[0]])
        TestNorm1 = np.linalg.norm(SVA_1[0:Sizes_C[0]])
        TestNorm2 = np.linalg.norm(SVA_2[0:Sizes_C[0]])

        print('\nTestNorm of first filter from core: \t\t{}'.format(TestNorm0))
        print('TestNorm of first filter from PCA vector 1: \t{}'.format(TestNorm1))
        print('TestNorm of first filter from PCA vector 2: \t{}'.format(TestNorm2))

        for x in range(len(Sizes_C) - 1):

            Seg0 = TrainedNetVector[Sizes_C[x]:Sizes_C[x+1]]
            Seg0_Norm = np.linalg.norm(Seg0)

            Seg1 = SVA_1[Sizes_C[x]:Sizes_C[x+1]]
            Seg1_Norm = np.linalg.norm(Seg1)

            Seg2 = SVA_2[Sizes_C[x]:Sizes_C[x+1]]
            Seg2_Norm = np.linalg.norm(Seg2)

            SVA_1[Sizes_C[x]:Sizes_C[x+1]] = SVA_1[Sizes_C[x]:Sizes_C[x+1]] * Seg0_Norm / Seg1_Norm
            SVA_2[Sizes_C[x]:Sizes_C[x+1]] = SVA_2[Sizes_C[x]:Sizes_C[x+1]] * Seg0_Norm / Seg2_Norm

    RIstep = int(round(FilterSteps/2))

    loss_array = np.zeros([FilterSteps+1,FilterSteps+1])

    for i in range(-RIstep, RIstep+1):
        for j in range(-RIstep, RIstep+1):
            NetAdd_1 = distance_multiplier*(i/RIstep)*SVA_1
            NetAdd_2 = distance_multiplier*(j/RIstep)*SVA_2
            NetWeights = TrainedNetVector + NetAdd_1 + NetAdd_2

            net_updated = list_to_weights(NetWeights, net, True)

            if CUDA_flag:
                net_updated.cuda()

            val_running_loss = 0.0
            val_loader = batch_generator (BS, val_set)

            with torch.no_grad():
                for val_inputs, val_labels in val_loader:

                    val_inputs = torch.stack(val_inputs)
                    val_inputs = torch.unsqueeze(val_inputs, 1)
                    val_labels = string_tensor(val_labels)
                    val_outputs = net_updated(val_inputs)
                    val_loss = criterion(val_outputs, val_labels.long())

                    val_running_loss += val_loss.item()

            loss_array[i+RIstep, j+RIstep] = val_running_loss

    x = np.arange(-RIstep+1, RIstep+1)
    y = np.arange(-RIstep, RIstep+1)
    X,Y = np.meshgrid(x,y)

    loss_array_cor = loss_array[:,1:]

    if bowl:

        tv0 = loss_array[round(FilterSteps/2) - 1, FilterSteps - 1]
        tv1 = loss_array[round(FilterSteps/2) - 1, 0]
        tv2 = loss_array[0, round(FilterSteps/2) - 1]
        tv3 = loss_array[FilterSteps - 1, round(FilterSteps/2) - 1]

        bowlmax = max([tv0, tv1, tv2, tv3])
        bowlmin = min([tv0, tv1, tv2, tv3])

        for vi in range(FilterSteps):
            for vj in range(FilterSteps):
                if loss_array[vi,vj] > bowlmin:
                    loss_array[vi,vj] = bowlmin

    PCA_fig_1 = plt.figure(figsize=(10,7))
    ax = PCA_fig_1.gca(projection='3d')
    PCA_surf = ax.plot_surface(X,Y, loss_array_cor, cmap='coolwarm')
    ax.set_xlabel(r'X', fontsize=20)
    ax.set_ylabel(r'Y', fontsize=20)
    ax.set_zlabel(r'Z', fontsize=20)

    PCA_fig_2 = plt.figure(figsize=(10,7))
    ax = PCA_fig_2.gca(projection='3d')
    PCA_surf_2 = ax.plot_surface(X,Y, loss_array_cor, cmap='coolwarm')
    ax.set_xlabel(r'X', fontsize=20)
    ax.set_ylabel(r'Y', fontsize=20)
    ax.set_zlabel(r'Z', fontsize=20)
    ax.view_init(30, 45)

    PCA_fig_3 = plt.figure(figsize=(10,7))
    PCA_cont = plt.contour(X,Y, loss_array_cor, 100)
    ax = PCA_fig_3.gca()
    ax.set_xlabel(r'X', fontsize=20)
    ax.set_ylabel(r'Y', fontsize=20)
    if contour_center:
        hlinex = [-19,20]
        hliney = [0,0]
        plt.plot(hlinex, hliney,'k')
        vlinex = [0,0]
        vliney = [-20,20]
        plt.plot(vlinex, vliney,'k')

        loss_array_cor[0][0] = loss_array_cor[0][1]
        LMIN = np.min(loss_array_cor)
        LMINLOC = np.argmin(loss_array_cor)
        RowLen = len(loss_array_cor[0])
        min_remain = LMINLOC % RowLen
        min_y = (LMINLOC - min_remain) / RowLen
        min_x = min_remain
        plt.plot(min_x - 20, min_y - 20, 'rx', markersize=15)

    PCA_time_end = time.time()
    PCA_time = PCA_time_end - PCA_time_start
    print('\nPCA time: {} s'.format(round(PCA_time,1)))

# %% ----- Saving results -----------------------------------------------------------------------------------------------

os.chdir(fname)

if MAV_flag:
    mavpath1 = '{}\\MAV.mat'.format(fname)
    mavpath2 = '{}\\MAV_dist_cor.mat'.format(fname)
    mavpath3 = '{}\\MAV_dist_inc.mat'.format(fname)
    sio.savemat(mavpath1, {'data':MAV})
    sio.savemat(mavpath2, {'data':AV_dist_correct})
    sio.savemat(mavpath3, {'data':AV_dist_wrong})

    f = open('AV stats.txt','w+')
    f.write('Correct predictions:')
    f.write('\n\nAV_max = {:.2f}'.format(AVC_max))
    f.write('\nAV_min = {:.2f}'.format(AVC_min))
    f.write('\nAV_avg = {:.2f}'.format(AVC_avg))
    f.write('\nAV_std = {:.2f}'.format(AVC_std))

    f.write('\n\nWrong predictions:')
    f.write('\n\nAV_max = {:.2f}'.format(AVW_max))
    f.write('\nAV_min = {:.2f}'.format(AVW_min))
    f.write('\nAV_avg = {:.2f}'.format(AVW_avg))
    f.write('\nAV_std = {:.2f}'.format(AVW_std))
    f.close()

f = open('run info.txt','w+')
A = str(net)
f.write('-----------------------------------------------------------------------------------------\n\n')
f.write(A)
f.write('\n\n-----------------------------------------------------------------------------------------')
f.write('\n\nData used:\n')

for i in range(len(N)):
    f.write('\nClass {}: {} images'.format(names[i], N[i]))

f.write('\n\nTrain set size: {} images'.format(len(train_set[0])))
f.write('\nVal set length: {} images'.format(len(val_set[0])))
f.write('\nTest set size: \t{} images'.format(len(test_set[0])))
f.write('\n\nTotal size of dataset: {} images'.format(sum(N)))

f.write('\n\n-----------------------------------------------------------------------------------------')

f.write('\n\ninput_size = {}'.format(input_size))
f.write('\nBS = {}'.format(BS))
f.write('\nepochs = {}'.format(epochs))
f.write('\nSTEPS = {}'.format(STEPS))
f.write('\nSeedVal = {}'.format(SeedVal))

f.write('\n\nLoss function: {}'.format(criterion))
f.write('\n\nOptimizer: \n{}'.format(optimizer))

f.write('\n\n-----------------------------------------------------------------------------------------')

f.write('\n\nAccuracy of the network on the test images: {:.1f} %'.format((100 * correct[0] / total[0])))
f.write('\nTest result obtained from {} images coming from {} batches'.format(total[0], testsize[0]))
f.write('\n\nAccuracy of the network on the total set of all images: {:.1f} %'.format((100 * correct_tot / total_tot)))
f.write('\nTest result obtained from {} images coming from {} batches'.format(total_tot, testsize_tot))

if MAV_flag:
    f.write('\n\nAmount of images in high confidence set: {}'.format(ImagesChecked))
    f.write('\nThis is {:.1f}% of the test set'.format(100*ImagesChecked / total[0]))

    f.write('\n\nAmount of correctly classified images in high confidence set: {}'.format(ImagesCheckedCorrect))
    f.write('\nAccuracy on high confidence set: {:.1f} %'.format(100*ImagesCheckedCorrect / ImagesChecked))

if PCA_flag:
    f.write('\n\n-----------------------------------------------------------------------------------------')
    f.write('\n\nPCA - filter normalization: {}'.format(filter_normalization_flag))
    f.write('\nPCA - distance multiplier: {}'.format(distance_multiplier))
    f.write('\nPCA - FilterSteps: {}'.format(FilterSteps))

if misclassified_outputs_flag:
    f.write('\n\n-----------------------------------------------------------------------------------------')
    f.write('\n\nMisclassified net outputs: ')
    f.write('\n')
    maxdif_t2 = [round(x,1) for x in maxdif_t]
    for x in range(len(misc_outputs_t)):
        f.write('\n{} \t- maxdif = {}'.format(misc_outputs_t[x], maxdif_t2[x]))
f.write('\n\n-----------------------------------------------------------------------------------------')

f.close()

cwd = os.getcwd()

torch.save({'model_state_dict': net.state_dict()
            }, r'{}\\final_weights.pth'.format(cwd))

torch.save({'model_state_dict': model_initial.state_dict()
            }, r'{}\\initial_weights.pth'.format(cwd))

path  = '{}\\loss_curve.png'.format(cwd)
val_plot.savefig(path, dpi=fig_dpi)
path  = '{}\\accuracy_curve.png'.format(cwd)
acc_plot.savefig(path, dpi=fig_dpi)
path  = '{}\\confusion_matrix.png'.format(cwd)
conf_fig.savefig(path, dpi=fig_dpi)
path  = '{}\\confusion_matrix_normalize_row.png'.format(cwd)
conf_fig_2.savefig(path, dpi=fig_dpi)
path  = '{}\\confusion_matrix_normalize_full.png'.format(cwd)
conf_fig_3.savefig(path, dpi=fig_dpi)

if PCA_flag:
    path  = '{}\\PCA_loss_landscape_1.png'.format(cwd)
    PCA_fig_1.savefig(path, dpi=fig_dpi)
    path  = '{}\\PCA_loss_landscape_2.png'.format(cwd)
    PCA_fig_2.savefig(path, dpi=fig_dpi)
    path  = '{}\\PCA_loss_contour.png'.format(cwd)
    PCA_fig_3.savefig(path, dpi=fig_dpi)

if loss_landscape_flag:
    path  = '{}\\loss_surface_contour.png'.format(cwd)
    loss_con.savefig(path, dpi=fig_dpi)
    path  = '{}\\loss_surface_angle_1.png'.format(cwd)
    loss_surf_1.savefig(path, dpi=fig_dpi)
    path  = '{}\\loss_surface_angle_2.png'.format(cwd)
    loss_surf_2.savefig(path, dpi=fig_dpi)

if MAV_flag:
    path = '{}\\MAV_distance_dist.png'.format(cwd)
    MAV_dist_plot.savefig(path, dpi=fig_dpi)

if misclassified_images_flag:
    mcfname = 'misclassified images'
    os.mkdir(mcfname)
    mcfdir = '{}\\{}'.format(fname, mcfname)
    os.chdir(mcfdir)

    for i in range(len(misc_preds_t)):

        imagets0 = misc_img_t[i][0]
        imagets = imagets0.cpu()
        predts = int(misc_preds_t[i])
        predname = names[predts]
        labts = int(misc_labels_t[i])
        labname = names[labts]
        fstag = '{} - label {}, classified as {}'.format(i, labname, predname)

        sf = plt.figure(figsize=(6,6))
        plt.imshow(imagets, cmap='jet')
        path = '{}\\{}.png'.format(mcfdir, fstag)
        sf.savefig(path)

if baseline_flag:
    path = '{}\\baseline_uncertainty_plot.png'.format(cwd)
    baseline_plot.savefig(path, dpi=fig_dpi)

    sio.savemat('unc_correct.mat', {'data':unc_correct})
    sio.savemat('unc_wrong.mat', {'data':unc_wrong})

os.chdir(DT)
#os.startfile(fname)


import os
import requests
import re
import urllib

directory = '../saved_page/photos'

# Create directory to save results
if not os.path.exists(directory):
    os.makedirs(directory)

def get_albums_list(user_id, secret_key):
    print("Getting photo albums...")
    photos_albums_response = requests.get("https://api.vk.com/method/photos.getAlbums?owner_id=" + user_id +
                                  "&need_system=1&access_token="+ secret_key).json()
    #print "Got {0} albums from UID {1}.".format(len(music_response[u'response'])-1, user_id)
    albums_list = []
    rx = '[' + re.escape('><\/:|?"*') + ']'
    for album in photos_albums_response[u'response']:
        # Remove characters that are forbidden in Windows filenames
        title = re.sub(rx, '', album[u'title'])
        folder = '{0}/{1}'.format(directory, title.encode('cp1251'))
        if not os.path.exists(folder):
            os.makedirs(folder)
        albums_list.append([title, album[u'aid']])
    print("Got {0} photo albums.".format(len(albums_list)))
    return albums_list

def get_photos_from_album(albums_list, user_id, secret_key):
    print("Getting photos...")
    c = 1
    sizes = [u'src',u'src_big',u'src_xbig', u'src_xxbig',u'src_xxxbig']
    for album in albums_list:
        print("{0}/{1}. Downloading {2}...".format(c,len(albums_list), album[0].encode("utf-8")))
        photos_response = requests.get("https://api.vk.com/method/photos.get?owner_id=" + user_id +
                                  "&album_id="+ str(album[1]) +"&access_token="+ secret_key).json()[u'response']
        for photo in photos_response:
            # Choosing the largest image available
            for size in sizes:
                try:
                    url = photo[size]
                except:
                    break
            picture = "{0}/{1}/{2}.jpg".format(directory, album[0].strip().encode("cp1251"),
                                                                      photo[u'pid'])
            if not os.path.isfile(picture):
                urllib.URLopener().retrieve(url, picture)
        c+=1

def get_photos(user_id, secret_key):
    albums_list = get_albums_list(user_id, secret_key)
    get_photos_from_album(albums_list, user_id, secret_key)


from django.urls import path, include
from rest_framework.routers import DefaultRouter
from profiles_api import views

#viewset
#configuring url to point to viewset; use a router; different urls for different methods unlike apiviewset
router = DefaultRouter()
router.register('hello-viewset', views.HelloViewSet, base_name='hello-viewset') #basename to retrieve url
router.register('profile', views.UserProfileViewSet) #dont need a basename cause we have queryset in view
#apiview set
router.register('feed', views.UserProfileFeedViewSet)
urlpatterns=[
	path('hello-view', views.HelloApiView.as_view()),
	path('login/', views.UserLoginApiView.as_view()),
	path('', include(router.urls))
]