from Csmake.CsmakeModule import CsmakeModule
from CsmakeProviders.GitProvider import GitProvider
import urlparse
from datetime import datetime
import sys
import yaml
import os
import os.path

class CsversionGenerateChangeLog(CsmakeModule):
    """Purpose: Create a change log from the mapping provided.
                Unless otherwise specified, the change log will
                be between the current state and the manifest
                provided in a mapping.  See "Options" for details
                This module should run after product metadata is
                finished.

    Options: previous-group - (OPTIONAL) This will define the group
                 to pull the manifest from.
                 (the group being the group id of the manifest to use
                  from the file mapping.  If multiple files match
                  the manifests will be grouped together allowing for
                  partial manifests to be generated into a single changelog)
                 Default is all files provided in mapping will be
                  used.  If a 1-1 mapping is specified, this can
                  produce several change logs at one time.
             current-group - (OPTIONAL) This will define the group
                 to compare the manifest "previous" against.
                 (the group being the group od of the manifest to use
                  from the file mapping, one will be picked.)
                 Default is to use the in memory __Csversion__ information
             git-stash - (OPTIONAL) If there is a location for git repos
                 that can be used to speed up the git delta generation
                 then point 'git-stash' to that location.
                 Default is %(RESULTS)s
             find-refs - (OPTIONAL) This will add related references
                 to the output when True - this may slow the build down a lot,
                 but will provide all gerrit review tags, for example,
                 which could be extremely useful
                 Default is False

    Mapping: Use **maps to map manifests into a change log on the right.
             Only files with a type of yaml and a purpose of manifest
             are accepted as inputs.
             Output will be based on the type(s) given for the output
                files.  Currently supported are 'yaml' and 'log'
                a file type that does not match the supported types
                will output a yaml file.
    Uses Environment:
        __Csversion__ - A dictionary where product metadata is stored under
                        'product'.  'product' is product info keyed off of
                        the type of data stored, i.e., metadata and build.
                        The same metadata/tag combination will be overwritten
                        if pulled twice.
                        The structure of the product dictionary is
                        a dictionary of tags from builds with metadata and
                        build information.
        __DIBEnv__ - If available, the git fetch will use the environment
                     from the DIB shell to maintain consistency within the build
                     When not available, the execution shell environment will
                     be used.
        """

    def _formatOutput(self, changelog, tofiles):
        class _TypeSetter:
            #TODO: Consider pulling this out so that other report types
            #      could be supported by subclassing
            def yaml(innerself, spec):
                yamltext = yaml.safe_dump(changelog)
                with open(spec['location'],'w') as log:
                    log.write(yamltext)

            def log(innerself, spec):
                self.log.error("TODO: create a human readable changelog file - requested but not implemented")
                #Do fancy stuff that Shelly would rather read
                #TODO: Lots of fancy formatting...

        ts = _TypeSetter()
        for spec in tofiles:
            #TODO: Consider adding type__intent
            if hasattr(ts, spec['type']):
                getattr(ts, spec['type'])(spec)
            else:
                getattr(ts, 'yaml')(spec)

    def _analyzeSpecificChange(self, packageType, records, build):
        self.log.devdebug("records: %s", str(records))
        old = records['old'][build][1]
        new = records['new'][build][1]
        class _PackageAnalyzer:
            def git(innerself):
                info = "Build: %s" % build
                delta = False

                url = None
                newref = None
                oldref = None
                if 'REF' not in old or 'REF' not in new:
                    info += ", a manifest was missing the git SHA1" % build
                    delta = True
                elif old['REF'] != new['REF']:
                    if 'URL' in new:
                        url = new['URL']
                    elif 'URL' in old:
                        url = old['URL']
                    if url is None:
                        info += ", there is no URL to access the git repo"
                    delta = True
                    newref = new['REF']
                    oldref = old['REF']
                else:
                    info += ", SHA1's are the same"

                if 'BRANCH' not in old or 'BRANCH' not in new \
                    or old['BRANCH'] != new['BRANCH']:
                    info += ", The branch changed from '%s' to '%s'" % (
                        old['BRANCH'],
                        new['BRANCH'] )
                    delta = True

                #get the sha delta from GitProvider
                if delta and url is not None:
                    _, name = os.path.split(url)
                    localrepo = os.path.join(
                        self.gitstash,
                        name )
                    result = None
                    try:
                        result = GitProvider.generateSpanLogInfo(
                            self.log,
                            localrepo,
                            oldref,
                            newref,
                            self.findrefs )
                    except Exception as e:
                        self.log.info("Span log threw '%s' attempting fetch and trying again", str(e))
                        try:
                            GitProvider.fetchRepository(
                                self.log,
                                name,
                                localrepo,
                                url,
                                newref,
                                None,
                                env=self.shellenv,
                                secure=False )
                            result = GitProvider.generateSpanLogInfo(
                                self.log,
                                localrepo,
                                oldref,
                                newref,
                                self.findrefs )
                        except:
                            self.log.exception("Attempt to generate git changelog failed")
                            info += ", Could not get git changelog due to failures"
                    if result is not None:
                        if 'changelog' not in records:
                            records['changelog'] = {}
                        records['changelog'][build] = result

                if delta:
                    records['__info'].append(info)
                return delta

            def dpkg(innerself):
                delta = 'VERSION' in old and 'VERSION' in new \
                       and old['VERSION'] != new['VERSION']
                if delta:
                    records['__info'].append(
                        "Build: %s, version changed from '%s' to '%s'" % (
                            build,
                            old['VERSION'],
                            new['VERSION'] ) )
                return delta

            def pip(innerself):
                delta = 'VERSION' in old and 'VERSION' in new \
                       and old['VERSION'] != new['VERSION']
                if delta:
                    records['__info'].append(
                        "Build: %s, version changed from '%s' to '%s'" % (
                            build,
                            old['VERSION'],
                            new['VERSION'] ) )
                return delta

            def file(innerself):
                info = "Build: '%s'" % build
                delta = False
                if 'INJECTED' not in old or 'INJECTED' not in new \
                       or old['INJECTED'] != new['INJECTED']:
                    info += ", image location changed"
                    delta = True
                if 'SHA1' not in old or 'SHA1' not in new \
                       or old['SHA1'] != new['SHA1']:
                    info += ", file content changed"
                    delta = True
                if 'URL' not in old or 'URL' not in new \
                       or old['URL'] != new['URL']:
                    info += ", pull url changed"
                    delta = True
                if delta:
                    records['__info'].append(info)
                return delta

            def tar(innerself):
                info = "Build: '%s'" % build
                delta = False
                if 'FILES' not in old or 'FILES' not in new \
                        or old['FILES'][0] != new['FILES'][0]:
                    info += ", image location changed"
                    delta = True
                if 'SHA1' not in old or 'SHA1' not in new \
                        or old['SHA1'] != new['SHA1']:
                    info += ", tar content changed"
                    delta = True
                if 'URL' not in old or 'URL' not in new \
                        or old['URL'] != new['URL']:
                    info += ", pull url changed"
                    delta = True
                if delta:
                    records['__info'].append(info)
                return delta

        pa = _PackageAnalyzer()
        if hasattr(pa, packageType):
            if not getattr(pa, packageType)():
                #records['__info'].append("Build: %s, shows no actual diff (TODO: remove from changelog)" % build)
                #TODO: Analysis shows no real change - delete the build record
                del records['old'][build]
                del records['new'][build]

        return len(records['old']) != 0 or len(records['new']) != 0

    def _createProductDiff(self, changelog, previous, current):
        if 'product' not in previous or 'product' not in current:
            return
        changelog['product']={
            'build': {
                'new' : current['product']['build'],
                'old' : previous['product']['build'],
                '__info' : [] } }
        for build, values in current['product']['build'].iteritems():
            if build not in previous['product']['build']:
                changelog['product']['build']['__info'].append(
                    "Build %s, previous didn't have a command" % build )
                continue

            if values['command'] != \
               previous['product']['build'][build]['command']:
                changelog['product']['build']['__info'].append("NOTE: Build: %s Command line changed" % build)

        for build in previous['product']['build'].keys():
            if build not in current['product']['build']:
                changelog['product']['build']['__info'].append(
                    "Build %s, new build didn't have a command" % build )

        changelog['product']['metadata']={}
        changemeta = changelog['product']['metadata']
        for build, metadata in current['product']['metadata'].iteritems():
            if build not in previous['product']['metadata']:
                changemeta[build] = {
                    '__info' : ['build in current and not in previous'] }
                continue
            changemeta[build] = {}
            for key, value in metadata.iteritems():
                if key not in previous['product']['metadata'][build]:
                    changemeta[build][key] = {
                        'new':value,
                        'old':"" }
                elif previous['product']['metadata'][build][key] != value:
                    changemeta[build][key] = {
                        'new':value,
                        'old':previous['product']['metadata'][build][key] }
            for key, value in previous['product']['metadata'][build].iteritems():
                if key not in metadata:
                    changemeta[build][key] = {
                        'new':"",
                        'old':value }

        for build in previous['product']['metadata'].keys():
            if build not in current['product']['metadata']:
                changemeta[build] = {
                    '__info': ['build in previous and not in current'] }
        self.log.devdebug("Changelog after product: %s", str(changelog))

    def _createSourcesDiff(self, changelog, previous, current):
        if 'sources' not in previous or 'sources' not in current:
            return
        changesource = {}
        changelog['sources'] = changesource
        for source, builds in current['sources'].iteritems():
            if source not in previous['sources']:
                changesource[source] = {
                    'new':builds,
                    'old':{},
                    '__info': ['source was not in previous'] }
                continue
            for build, info in builds.iteritems():
                if build not in previous['sources'][source]:
                    if source not in changesource:
                        changesource[source] = {
                            'new':{},
                            'old':{},
                            '__info': [] }
                    changesource[source]['new'][build] = info
                    changesource[source]['old'][build] = {}
                    changesource[source]['__info'].append(
                        "Source was not in previous for build: '%s'" % build )
                else:
                    previnfo = previous['sources'][source][build]
                    if info[0] != previnfo[0]:
                        if source not in changesource:
                            changesource[source] = {
                                'new':{},
                                'old':{},
                                '__info': [] }
                        changesource[source]['new'][build] = info
                        changesource[source]['old'][build] = previnfo
                        changesource[source]['__info'].append(
                            "Source changed formats for build: '%s'" % build )
                    else:
                        delta = False
                        for key, value in info[1].iteritems():
                            self.log.devdebug("previnfo: %s", str(previnfo))
                            if key not in previnfo[1]:
                                if source not in changesource:
                                    changesource[source] = {
                                        'new':{},
                                        'old':{},
                                        '__info': [] }
                                changesource[source]['new'][build] = info
                                changesource[source]['old'][build] = previnfo
                                changesource[source]['__info'].append(
                                    "Format description changed for '%s' in build: '%s'" % (
                                        info[0],
                                        build ) )
                                delta = True
                            elif value != previnfo[1][key]:
                                if source not in changesource:
                                    changesource[source] = {
                                        'new':{},
                                        'old':{},
                                        '__info':[] }
                                changesource[source]['new'][build] = info
                                changesource[source]['old'][build] = previnfo
                                delta = True
                            if delta:
                                break
                        if not delta:
                            for key in previnfo[1].keys():
                                if key not in info[1]:
                                    if source not in changesource:
                                        changesource[source] = {
                                            'new':{},
                                            'old':{},
                                            '__info':[] }
                                    changesource[source]['new'][build] = info
                                    changesource[source]['old'][build] = previnfo
                                    changesource[source]['__info'].append(
                                        "Format description changed for '%s' in build: '%s'" % (
                                            info[0],
                                            build ) )
                                    delta = True
                    if delta:
                        #dispatch specific info[0] analysis (package specific)
                        if not self._analyzeSpecificChange(
                            info[0],
                            changesource[source],
                            build ):
                            del changesource[source]

    def build(self, options):
        if self.mapping is None:
            self.log.warning("**maps was not specified in section.  Nothing to do...")
            self.log.warning("   No changelog will be produced")
            self.log.passed()
            return None

        previousGroup = None
        if 'previous-group' in options:
            previousGroup = options['previous-group']

        currentGroup = None
        if 'current-group' in options:
            currentGroup = options['current-group']

        self.gitstash = self.env.env['RESULTS']
        if 'git-stash' in options:
            self.gitstash = options['git-stash']

        self.findrefs = False
        if 'find-refs' in options:
            self.findrefs = options['find-refs'] == 'True'

        self.shellenv = os.environ
        if '__DIBEnv__' in self.env.env:
            if 'shellenv' in self.env.env['__DIBEnv__']:
                self.shellenv = self.env.env['__DIBEnv__']['shellenv']

        groupings = []
        for (fromSpecs, toSpecs) in self.mapping.iterspecs():
            current = {'previous':[], 'current':[]}
            groupings.append(current)
            for fromSpec in fromSpecs:
                if previousGroup is None or previousGroup == fromSpec['id']:
                    current['previous'].append(fromSpec)
                if currentGroup is not None and currentGroup == fromSpec['id']:
                    current['current'].append(fromSpec)
            current['to'] = toSpecs

        for group in groupings:
            previous = {}
            for spec in group['previous']:
                with open(spec['location']) as f:
                    mani = yaml.load(f)
                previous.update(mani)
            if len(previous) == 0:
                self.log.warning("**maps did not yield any manifests to map")
                self.log.warning("   No changelog will be produced")
                self.log.passed()
                return None
            current = {}
            if currentGroup is not None:
                if len(group['current']) == 0:
                    self.log.warning(
                        "The specified current-group '%s' didn't have any files in the map",
                        currentGroup )
                    self.log.warning(
                        "    No changelog will be produced" )
                    self.log.passed()
                    return None
                for spec in group['current']:
                    with open(spec['location']) as f:
                        mani = yaml.load(f)
                    current.update(mani)
                if len(current) == 0:
                    self.log.warning(
                        "A current-group was specified but no manifests were found")
                    self.log.warning(
                        "    No changelog will be produced" )
                    self.log.passed()
                    return None

            else:
                if '__Csversion__' not in self.env.env:
                    self.log.warning(
                        "A current-group was not specified, and there is no internal build state to reference")
                    self.log.warning(
                        "    No changelog will be produced" )
                    return None
                current = self.env.env['__Csversion__']
            changelog = {}
            self._createProductDiff(changelog, previous, current)
            self._createSourcesDiff(changelog, previous, current)
            self._formatOutput(changelog, toSpecs)

        self.log.passed()
        return True


import rdflib
import sys
import os
#import time
from rdflib import *
#__requires__ = "owlready2==0.13"
#import pkg_resources
#import owlready2
from owlready2 import *

alignOnto1 = rdflib.term.URIRef('http://knowledgeweb.semanticweb.org/heterogeneity/alignmentonto1')
alignOnto2 = rdflib.term.URIRef('http://knowledgeweb.semanticweb.org/heterogeneity/alignmentonto2')
alignCell = rdflib.term.URIRef('http://knowledgeweb.semanticweb.org/heterogeneity/alignmentCell')
alignEntity1 = rdflib.term.URIRef('http://knowledgeweb.semanticweb.org/heterogeneity/alignmententity1')
alignEntity2 = rdflib.term.URIRef('http://knowledgeweb.semanticweb.org/heterogeneity/alignmententity2')
alignMeasure = rdflib.term.URIRef('http://knowledgeweb.semanticweb.org/heterogeneity/alignmentmeasure')
alignRelation = rdflib.term.URIRef('http://knowledgeweb.semanticweb.org/heterogeneity/alignmentrelation')

def descendants(e):
    r = set()
    r.add(e)
    toVisit = set()
    toVisit.add(e)
    while len(toVisit) > 0:
        next = set()
        for c in toVisit:
            for d in c.subclasses():
                if d not in r:
                    r.add(d)
                    next.add(d)
        toVisit = next
    return r

def correct(v,o1,o2):
    t0 = time.process_time()
    aligned1 = {}
    aligned2 = {}

    tx1 = {}
    tx2 = {}
    l = {}
    l1 = len(o1.base_iri)
    l2 = len(o2.base_iri)
    for s in v.subjects(RDF.type, alignCell):
        e1 = v.value(s,alignEntity1, None)[l1:]
        e2 = v.value(s,alignEntity2, None)[l2:]
        if e1 not in aligned1:
            aligned1[e1] = set()
        if e2 not in aligned2:
            aligned2[e2] = set()
        aligned1[e1].add(s)
        aligned2[e2].add(s)
        tx1[s] = {}
        tx2[s] = {}
        tx1[s][s] = True
        tx2[s][s] = True
        l[s]=(s,e1,e2)

    t1 = time.process_time()
    droplist = []
    for i in l:
        ln = []
        for y in o1[l[i][1]].ancestors():
            if y.name in aligned1:
                for s in aligned1[y.name]:
                    tx1[i][s] = True
                    ln.append(y.name)
        for y in o2[l[i][2]].ancestors():
            if y.name in aligned2:
                for s in aligned2[y.name]:
                    if not s in tx1[i]:
                        if i in aligned1[l[i][1]]:
                            aligned1[l[i][1]].remove(i)
                            aligned2[l[i][2]].remove(i)
                            droplist.append(i)
                            aligned1[l[s][1]].remove(s)
                            aligned2[l[s][2]].remove(s)
                            droplist.append(s)
                        break
                    else:
                        tx2[i][s] = True
                else:
                    continue
                break
        else:
            for y in ln:
                for s in aligned1[y]:
                    if not s in tx2[i]:
                        if i in aligned1[l[i][1]]:
                            aligned1[l[i][1]].remove(i)
                            aligned2[l[i][2]].remove(i)
                            droplist.append(i)
                            aligned1[l[s][1]].remove(s)
                            aligned2[l[s][2]].remove(s)
                            droplist.append(s)
                        break
                else:
                    continue
                break

    t2 = time.process_time()

    newdrop = droplist
    newdrop = []
    for i in droplist:
        ln = []
        for y in o1[l[i][1]].ancestors():
            if y.name in aligned1:
                for s in aligned1[y.name]:
                    tx1[i][s] = True
                    ln.append(y.name)
        ld = []
        for y in descendants(o1[l[i][1]]):
            if y.name in aligned1:
                for s in aligned1[y.name]:
                    tx1[s][i] = True
                    ld.append(y.name)
        for y in o2[l[i][2]].ancestors():
            if y.name in aligned2:
                for s in aligned2[y.name]:
                    if not s in tx1[i]:
                        newdrop.append(i)
                        break
                    else:
                        tx2[i][s] = True
                else:
                    continue
                break
        else:
            for y in descendants(o2[l[i][2]]):
                if y.name in aligned2:
                    for s in aligned2[y.name]:
                        if not i in tx1[s]:
                            newdrop.append(i)
                            break
                        else:
                            tx2[s][i] = True
                    else:
                        continue
                    break
            else:
                for y in ln:
                    for s in aligned1[y]:
                        if not s in tx2[i]:
                            newdrop.append(i)
                            break
                    else:
                        continue
                    break
                else:
                    for y in ld:
                        for s in aligned1[y]:
                            if not i in tx2[s]:
                                newdrop.append(i)
                                break
                        else:
                            continue
                        break
                    else:
                        aligned1[l[i][1]].add(i)
                        aligned2[l[i][2]].add(i)
    t3 = time.process_time()
    return (newdrop,(t1-t0,t2-t1,t3-t2, len(droplist),len(newdrop)))

def printAlignment(v,out):
    sys.stderr = open(os.devnull, "w")
    out.write('<?xml version="1.0" encoding="utf-8"?>\n')
    out.write('<rdf:RDF xmlns="http://knowledgeweb.semanticweb.org/heterogeneity/alignment"\n')
    out.write('\txmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"\n')
    out.write('\txmlns:xsd="http://www.w3.org/2001/XMLSchema#">\n<Alignment>\n<xml>yes</xml>\n<level>0</level>\n<type>??</type>\n')
    out.write('<onto1><Ontology rdf:about="'+list(v.triples((None,alignOnto1,None)))[0][2].lower()+'"></Ontology></onto1>\n')
    out.write('<onto2><Ontology rdf:about="'+list(v.triples((None,alignOnto2,None)))[0][2].lower()+'"></Ontology></onto2>\n')
    for s in v.subjects(RDF.type, alignCell):
        out.write('<map>\n\t<Cell>\n')
        out.write('\t\t<entity1 rdf:resource="'+v.value(s,alignEntity1, None)+'"/>\n')
        out.write('\t\t<entity2 rdf:resource="'+v.value(s,alignEntity2, None)+'"/>\n')
        out.write('\t\t<measure rdf:datatype="xsd:float">'+v.value(s,alignMeasure, None)+'</measure>\n')
        out.write('\t\t<relation>'+v.value(s,alignRelation, None)+'</relation>\n')
        out.write('\t</Cell>\n</map>\n')
    out.write('</Alignment>\n</rdf:RDF>\n')
    sys.stderr = sys.__stderr__

def evaluateAlignment(align,onto1,onto2):
    v = rdflib.Graph().parse(align)
    ogSize = len(list(v.subjects(RDF.type, alignCell)))

    o1 = get_ontology(onto1)
    o1.load()

    o2 = get_ontology(onto2)
    o2.load()

    sys.stderr = open(os.devnull, "w") # hide reasoner verbosity
    sync_reasoner_pellet()
    sys.stderr = sys.__stderr__

    (remove,data) = correct(v,o1,o2)

    for n in remove:
        v.remove((n, None, None))
        v.remove((None, None, n))

    o = open("corrected_"+align,'w')
    printAlignment(v,o)
    o.close()

    print("Correction complete!\n===Summary===\nInput alignment: %s\tSize: %d"%(align,ogSize))
    print("Output alignment: corrected_%s\tSize %d\nTime for repair computation (greedy algorithm): %.5f seconds"%(align,ogSize-data[4],data[1]))
    print("Time for repair improvement: %.5f seconds\nPreliminary repair size (greedy algorithm):%d\nFinal repair size (after improvement step): %d\n"%(data[2],data[3],data[4]))

if __name__ == "__main__":
    if len(sys.argv) == 4:
        evaluateAlignment(sys.argv[1], sys.argv[2], sys.argv[3])
    else:
        print("Incorrect number of parameters! Usage:\npython3 conservativity.py <alignment file> <ontology1 file> <ontology 2 file>")


#!/usr/bin/env python3
"""
Contains function to sum a list of integers and floats

function:
 * sum_mixed_list: sums up a list
"""

from typing import List, Union

def sum_mixed_list(mxd_lst: List[Union[int, float]]) -> float:
    """Sums up a list of integers and floats"""
    total_sum: float = 0.0
    for i in mxd_lst:
        total_sum += i
    return total_sum


from . import parcellationmap, volume as _volume

from ..commons import MapIndex, logger, iterate_connected_components, siibra_tqdm
from ..locations import boundingbox
from ..retrieval import cache
from ..retrieval.repositories import ZipfileConnector, GitlabConnector

from os import path, rename, makedirs
from zipfile import ZipFile, ZIP_DEFLATED
import gzip
from typing import Dict, Union, TYPE_CHECKING, List
from nilearn import image
from nibabel import Nifti1Image, load
import numpy as np

if TYPE_CHECKING:
    from ..core.region import Region

class SparseIndex:

    def __init__(self):
        self.probs = []
        self.bboxes = []

        # these are initialized when adding the first volume, see below
        self.affine: np.ndarray = None
        self.shape = None
        self.voxels: np.ndarray = None

    def add_img(self, img: Nifti1Image):

        if self.num_volumes == 0:
            self.affine = img.affine
            self.shape = img.shape
            self.voxels = np.zeros(img.shape, dtype=np.int32) - 1
        else:
            if (img.shape != self.shape) or ((img.affine - self.affine).sum() != 0):
                raise RuntimeError(
                    "Building sparse maps from volumes with different voxel spaces is not yet supported in siibra."
                )

        volume = self.num_volumes
        imgdata = np.asanyarray(img.dataobj)
        X, Y, Z = [v.astype("int32") for v in np.where(imgdata > 0)]
        for x, y, z, prob in zip(X, Y, Z, imgdata[X, Y, Z]):
            coord_id = self.voxels[x, y, z]
            if coord_id >= 0:
                # Coordinate already seen. Add observed value.
                assert volume not in self.probs[coord_id]
                assert len(self.probs) > coord_id
                self.probs[coord_id][volume] = prob
            else:
                # New coordinate. Append entry with observed value.
                coord_id = len(self.probs)
                self.voxels[x, y, z] = coord_id
                self.probs.append({volume: prob})

        self.bboxes.append(
            {
                "minpoint": (X.min(), Y.min(), Z.min()),
                "maxpoint": (X.max(), Y.max(), Z.max()),
            }
        )

    @property
    def num_volumes(self):
        return len(self.bboxes)

    def max(self):
        return self.voxels.max()

    def coords(self, volume: int):
        # Nx3 array with x/y/z coordinates of the N nonzero values of the given mapindex
        assert volume in range(self.num_volumes)
        coord_ids = [i for i, l in enumerate(self.probs) if volume in l]
        x0, y0, z0 = self.bboxes[volume]["minpoint"]
        x1, y1, z1 = self.bboxes[volume]["maxpoint"]
        return (
            np.array(
                np.where(
                    np.isin(
                        self.voxels[x0: x1 + 1, y0: y1 + 1, z0: z1 + 1],
                        coord_ids,
                    )
                )
            ).T
            + (x0, y0, z0)
        ).T

    def mapped_voxels(self, volume: int):
        # returns the x, y, and z coordinates of nonzero voxels for the map
        # with the given index, together with their corresponding values v.
        assert volume in range(self.num_volumes)
        x, y, z = [v.squeeze() for v in np.split(self.coords(volume), 3)]
        v = [self.probs[i][volume] for i in self.voxels[x, y, z]]
        return x, y, z, v

    def _to_local_cache(self, cache_prefix: str):
        """
        Serialize this index to the cache, using the given prefix for the cache
        filenames.
        """
        probsfile = cache.CACHE.build_filename(f"{cache_prefix}", suffix="probs.txt.gz")
        bboxfile = cache.CACHE.build_filename(f"{cache_prefix}", suffix="bboxes.txt.gz")
        voxelfile = cache.CACHE.build_filename(f"{cache_prefix}", suffix="voxels.nii.gz")
        Nifti1Image(self.voxels, self.affine).to_filename(voxelfile)
        with gzip.open(probsfile, 'wt') as f:
            for D in self.probs:
                f.write("{}\n".format(" ".join(f"{i} {p}" for i, p in D.items())))
        with gzip.open(bboxfile, "wt") as f:
            for bbox in self.bboxes:
                f.write(
                    "{} {}\n".format(
                        " ".join(map(str, bbox["minpoint"])),
                        " ".join(map(str, bbox["maxpoint"])),
                    )
                )

    @staticmethod
    def _from_local_cache(cache_name: str):
        """
        Attempts to build a sparse index from the siibra cache, looking for
        suitable cache files with the specified prefix.

        Parameters
        ----------
        prefix: str
            Prefix of the filenames.

        Returns
        -------
        SparseIndex
            None if cached files are not found or suitable.
        """

        probsfile = cache.CACHE.build_filename(f"{cache_name}", suffix="probs.txt.gz")
        bboxfile = cache.CACHE.build_filename(f"{cache_name}", suffix="bboxes.txt.gz")
        voxelfile = cache.CACHE.build_filename(f"{cache_name}", suffix="voxels.nii.gz")
        if not all(path.isfile(f) for f in [probsfile, bboxfile, voxelfile]):
            return None

        result = SparseIndex()

        voxels = load(voxelfile)
        result.voxels = np.asanyarray(voxels.dataobj)
        result.affine = voxels.affine
        result.shape = voxels.shape

        with gzip.open(probsfile, "rt") as f:
            lines = f.readlines()
            for line in siibra_tqdm(
                lines,
                total=len(lines),
                desc="Loading sparse index",
                unit="voxels"
            ):
                fields = line.strip().split(" ")
                mapindices = list(map(int, fields[0::2]))
                values = list(map(float, fields[1::2]))
                D = dict(zip(mapindices, values))
                result.probs.append(D)

        with gzip.open(bboxfile, "rt") as f:
            for line in f:
                fields = line.strip().split(" ")
                result.bboxes.append(
                    {
                        "minpoint": tuple(map(int, fields[:3])),
                        "maxpoint": tuple(map(int, fields[3:])),
                    }
                )

        return result

class SparseMap(parcellationmap.Map):
    """
    A sparse representation of list of statistical (e.g. probabilistic) brain
    region maps.

    It represents the 3D statistical maps of N brain regions by two data structures:

    1) 'spatial_index', a 3D volume where non-negative values represent unique indices into a list of region assignments
    2) 'probs', a list of region assignments where each entry is a dict

    More precisely, given ``i = sparse_index.voxels[x, y, z]`` we define that

    - if `i<0`, no brain region is assigned at this location
    - if `i>=0`, ``probs[i]`` defines the probabilities of brain regions.

    Each entry in probs is a dictionary that represents the region assignments for
    the unique voxel where ``spatial_index == i``. The assignment maps from a MapIndex
    to the actual (probability) value.
    """

    # A gitlab instance with holds precomputed sparse indices
    _GITLAB_SERVER = 'https://jugit.fz-juelich.de'
    _GITLAB_PROJECT = 5779

    def __init__(
        self,
        identifier: str,
        name: str,
        space_spec: dict,
        parcellation_spec: dict,
        indices: Dict[str, MapIndex],
        volumes: list = [],
        shortname: str = "",
        description: str = "",
        modality: str = None,
        publications: list = [],
        datasets: list = []
    ):
        parcellationmap.Map.__init__(
            self,
            identifier=identifier,
            name=name,
            space_spec=space_spec,
            parcellation_spec=parcellation_spec,
            indices=indices,
            shortname=shortname,
            description=description,
            modality=modality,
            publications=publications,
            datasets=datasets,
            volumes=volumes,
        )
        self._sparse_index_cached = None

    @property
    def _cache_prefix(self):
        return f"{self.parcellation.id}_{self.space.id}_{self.maptype}_{self.name}_index"

    @property
    def sparse_index(self):
        if self._sparse_index_cached is None:
            spind = SparseIndex._from_local_cache(self._cache_prefix)
            if spind is None:
                logger.info("Downloading precomputed SparseIndex...")
                gconn = GitlabConnector(self._GITLAB_SERVER, self._GITLAB_PROJECT, "main")
                zip_fname = f"{self.name.replace(' ', '_')}_index.zip"
                try:
                    assert zip_fname in gconn.search_files(), f"{zip_fname} is not in {gconn}."
                    zipfile = gconn.get_loader(zip_fname).url
                    spind = self.load_zipped_sparseindex(zipfile)
                except Exception:
                    logger.info("Failed to load precomputed SparseIndex from Gitlab.")
                    logger.debug(f"Could not load SparseIndex from Gitlab at {gconn}", exc_info=1)
            if spind is None:
                with _volume.SubvolumeProvider.UseCaching():
                    spind = SparseIndex()
                    for vol in siibra_tqdm(
                        range(len(self)), total=len(self), unit="maps",
                        desc=f"Fetching {len(self)} volumetric maps"
                    ):
                        img = super().fetch(
                            index=MapIndex(volume=vol, label=None)
                        )
                        if img is None:
                            region = self.get_region(volume=vol)
                            logger.error(f"Cannot retrieve volume #{vol} for {region.name}, it will not be included in the sparse map.")
                            continue
                        spind.add_img(img)
                    spind._to_local_cache(self._cache_prefix)
            self._sparse_index_cached = spind
        assert self._sparse_index_cached.max() == len(self._sparse_index_cached.probs) - 1
        return self._sparse_index_cached

    @property
    def affine(self):
        return self.sparse_index.affine

    @property
    def shape(self):
        return self.sparse_index.shape

    def save_sparseindex(self, destination: str, filename: str = None):
        """
        Save SparseIndex as a .zip in destination folder from local cache. If
        SparseIndex is not cached, siibra will firt create it first.

        Parameters
        ----------
        destination: str
            The path where the zip file will be created.
        filename: str, default=None
            Name of the zip and prefix of the SparseIndex files. If None, siibra
            uses `name` property.
        """
        if filename is None:
            filename = f"{self.name.replace(' ', '_')}_index"
        logger.info(f"Saving SparseIndex of '{self.name}' as '{filename}.zip'")
        if not path.isdir(destination):
            makedirs(destination)
        if self._sparse_index_cached is None:
            _ = self.sparse_index
        suffices = [".probs.txt.gz", ".bboxes.txt.gz", ".voxels.nii.gz"]
        try:
            with ZipFile(f"{destination}/{filename}.zip", 'w') as zipf:
                for suffix in suffices:
                    zipf.write(
                        filename=cache.CACHE.build_filename(self._cache_prefix, suffix),
                        arcname=path.basename(f"{filename}{suffix}"),
                        compress_type=ZIP_DEFLATED
                    )
        except Exception as e:
            logger.error("Could not save SparseIndex:\n")
            raise e
        logger.info("SparseIndex is saved.")

    def load_zipped_sparseindex(self, zipfname: str):
        """
        Load SparseIndex from previously computed source and creates a local
        cache.

        Parameters
        ----------
        zipfile: str
            A url or a path to zip file containing the SparseIndex files for
            this SparseMap precomputed by siibra.

        Returns
        -------
        SparseIndex
        """
        zconn = ZipfileConnector(zipfname)
        with ZipFile(zconn.zipfile, 'r') as zp:
            suffices = [".probs.txt.gz", ".bboxes.txt.gz", ".voxels.nii.gz"]
            for suffix in suffices:
                file = [f for f in zconn.search_files(suffix=suffix)]
                assert len(file) == 1, f"Could not find a unique '{suffix}' file in {zipfname}."
                zp.extract(file[0], cache.CACHE.folder)
                rename(
                    path.join(cache.CACHE.folder, file[0]),
                    cache.CACHE.build_filename(self._cache_prefix, suffix=suffix)
                )
        zconn.clear_cache()

        return SparseIndex._from_local_cache(self._cache_prefix)

    def fetch(
        self,
        region_or_index: Union[MapIndex, str, 'Region'] = None,
        *,
        index: MapIndex = None,
        region: Union[str, 'Region'] = None,
        **kwargs
    ):
        """
        Recreate a particular volumetric map from the sparse
        representation.

        Parameters
        ----------
        region_or_index: str, Region, MapIndex
            Lazy match the specification.
        index : MapIndex
            The index to be fetched.
        region: str, Region
            Region name specification. If given, will be used to decode the map
            index of a particular region.

        Returns
        -------
        An image or mesh
        """
        if kwargs.get('format') in ['mesh'] + _volume.Volume.MESH_FORMATS:
            # a mesh is requested, this is not handled by the sparse map
            return super().fetch(region_or_index, index=index, region=region, **kwargs)

        try:
            length = len([arg for arg in [region_or_index, region, index] if arg is not None])
            assert length == 1
        except AssertionError:
            if length > 1:
                raise parcellationmap.ExcessiveArgumentException(
                    "One and only one of region_or_index, region, index can be defined for fetch"
                )
            # user can provide no arguments, which assumes one and only one volume present

        if isinstance(region_or_index, MapIndex):
            index = region_or_index

        from ..core.region import Region
        if isinstance(region_or_index, (str, Region)):
            region = region_or_index

        volidx = None
        if index is not None:
            assert isinstance(index, MapIndex)
            volidx = index.volume
        if region is not None:
            index = self.get_index(region)
            assert index is not None
            volidx = index.volume

        if volidx is None:
            try:
                assert len(self) == 1
                volidx = 0
            except AssertionError:
                raise parcellationmap.InsufficientArgumentException(
                    f"{self.__class__.__name__} provides {len(self)} volumes. "
                    "Specify 'region' or 'index' for fetch() to identify one."
                )

        assert isinstance(volidx, int)
        x, y, z, v = self.sparse_index.mapped_voxels(volidx)
        result = np.zeros(self.shape, dtype=np.float32)
        result[x, y, z] = v
        return Nifti1Image(result, self.affine)

    def _read_voxel(self, x, y, z):
        spind = self.sparse_index
        vx = spind.voxels[x, y, z]
        if isinstance(vx, int):
            return list(
                (None, volume, None, value)
                for volume, value in spind.probs[vx].items()
            )
        else:
            return list(
                (pointindex, volume, None, value)
                for pointindex, voxel in enumerate(vx)
                for volume, value in spind.probs[voxel].items()
            )

    def _assign_image(self, queryimg: Nifti1Image, minsize_voxel: int, lower_threshold: float, split_components: bool = True) -> List[parcellationmap.AssignImageResult]:
        """
        Assign an image volume to this parcellation map.

        Parameters:
        -----------
        queryimg: Nifti1Image
            the image to be compared with maps
        minsize_voxel: int, default: 1
            Minimum voxel size of image components to be taken into account.
        lower_threshold: float, default: 0
            Lower threshold on values in the statistical map. Values smaller than
            this threshold will be excluded from the assignment computation.
        """
        assignments = []

        # resample query image into this image's voxel space, if required
        if (queryimg.affine - self.affine).sum() == 0:
            queryimg = queryimg
        else:
            if issubclass(np.asanyarray(queryimg.dataobj).dtype.type, np.integer):
                interp = "nearest"
            else:
                interp = "linear"
            queryimg = image.resample_img(
                queryimg,
                target_affine=self.affine,
                target_shape=self.shape,
                interpolation=interp,
            )

        querydata = np.asanyarray(queryimg.dataobj).squeeze()

        iter_func = iterate_connected_components if split_components \
            else lambda img: [(1, img)]

        for mode, modeimg in iter_func(queryimg):

            # determine bounding box of the mode
            modemask = np.asanyarray(modeimg.dataobj)
            XYZ2 = np.array(np.where(modemask)).T
            position = np.dot(modeimg.affine, np.r_[XYZ2.mean(0), 1])[:3]
            if XYZ2.shape[0] <= minsize_voxel:
                continue
            X2, Y2, Z2 = [v.squeeze() for v in np.split(XYZ2, 3, axis=1)]

            bbox2 = boundingbox.BoundingBox(XYZ2.min(0), XYZ2.max(0) + 1, space=None)
            if bbox2.volume == 0:
                continue

            spind = self.sparse_index

            for volume in siibra_tqdm(
                range(len(self)),
                desc=f"Assigning structure #{mode} to {len(self)} sparse maps",
                total=len(self),
                unit=" map"
            ):
                bbox1 = boundingbox.BoundingBox(
                    self.sparse_index.bboxes[volume]["minpoint"],
                    self.sparse_index.bboxes[volume]["maxpoint"],
                    space=None,
                )
                if bbox1.intersection(bbox2) is None:
                    continue

                # compute union of voxel space bounding boxes
                bbox = bbox1.union(bbox2)
                bbshape = np.array(bbox.shape, dtype="int") + 1
                x0, y0, z0 = map(int, bbox.minpoint)

                # build flattened vector of map values
                v1 = np.zeros(np.prod(bbshape))
                XYZ1 = spind.coords(volume).T
                X1, Y1, Z1 = [v.squeeze() for v in np.split(XYZ1, 3, axis=1)]
                indices1 = np.ravel_multi_index(
                    (X1 - x0, Y1 - y0, Z1 - z0), bbshape
                )
                v1[indices1] = [spind.probs[i][volume] for i in spind.voxels[X1, Y1, Z1]]
                v1[v1 < lower_threshold] = 0

                # build flattened vector of input image mode
                v2 = np.zeros(np.prod(bbshape))
                indices2 = np.ravel_multi_index(
                    (X2 - x0, Y2 - y0, Z2 - z0), bbshape
                )
                v2[indices2] = querydata[X2, Y2, Z2]

                assert v1.shape == v2.shape

                intersection = np.sum(
                    (v1 > 0) & (v2 > 0)
                )  # np.minimum(v1, v2).sum()
                if intersection == 0:
                    continue
                iou = intersection / np.sum(
                    (v1 > 0) | (v2 > 0)
                )  # np.maximum(v1, v2).sum()

                v1d = v1 - v1.mean()
                v2d = v2 - v2.mean()
                rho = (
                    (v1d * v2d).sum()
                    / np.sqrt((v1d ** 2).sum())
                    / np.sqrt((v2d ** 2).sum())
                )

                maxval = v1.max()

                assignments.append(
                    parcellationmap.AssignImageResult(
                        input_structure=mode,
                        centroid=tuple(position.round(2)),
                        volume=volume,
                        fragment=None,
                        map_value=maxval,
                        intersection_over_union=iou,
                        intersection_over_first=intersection / (v1 > 0).sum(),
                        intersection_over_second=intersection / (v2 > 0).sum(),
                        correlation=rho,
                        weighted_mean_of_first=np.sum(v1 * v2) / np.sum(v2),
                        weighted_mean_of_second=np.sum(v1 * v2) / np.sum(v1)
                    )
                )

        return assignments


import time
import torch
import torch.utils.data
import torch.optim as optim
from loss import *
from torch.autograd import Variable
from model import model
from Evaluation import evaluation
from torchvision import transforms
from sklearn.datasets import fetch_lfw_people
from utils import GenIdx,IdentitySampler, prepare_dataset, prepare_set, \
                    prepare_data_ids,LFW_training_Data
import numpy as np

#### Fonction utiles ####

# Fonction d'entraînement
def train(epoch, criterion, optimizer, trainloader, device, net):

    """

    Se réferer éventuellement au laboratoire 5
    On pensera à mettre les données sur 3 canaux au lieux de 1 via expand :
    https://pytorch.org/docs/stable/generated/torch.Tensor.expand.html

    """

# Fonction permettant la validation
def valid(query_gall_loader, query_gallery_labels):

    # Extraire les features pour chaque image de la base de données
    query_gall_feat_pool, query_gall_feat = extract_query_gall_feat(query_gall_loader, n_query_gall, net, test_batch_size)

    # Compute the similarity matrix based on the extracted features matrix:
    similarity_matrix_pool = "A définir"
    similarity_matrix_feat = "A définir"

    # Evaluation - Compute metrics (Rank-1, Rank-5, mAP)
    cmc_att, mAP_att  = evaluation(-similarity_matrix_pool, query_gallery_labels)
    cmc, mAP = evaluation(-similarity_matrix_feat, query_gallery_labels)

    return cmc, mAP, cmc_att, mAP_att

#
def extract_query_gall_feat(query_gall_loader, ngall, net, batch_size):
    """
    Fonction d'extraction des caractéristiques et création d'une matrice contenant sur chaque ligne
    le vecteur caractéristique d'une image.

    Deux matrices sont en fait produites :
        - L'une à partir du vecteur caractéristique produit après la couche average pooling
        - L'autre à partir du vecteur caractéristique produit après la couche batch normalization

    Dans la pratique, il est parfois plus intéressant de regarder les résultats de l'une ou de l'autre, donc
    on testera les deux.

    """
    net.eval()
    print('Extracting Gallery Feature...')
    # start = time.time()
    ptr = 0
    feat_size = 512
    gall_feat_pool = np.zeros((ngall, feat_size))
    gall_feat = np.zeros((ngall, feat_size))

    with torch.no_grad():
        # Boucle itérant sur chaque batch de validation à définir ici
        for A_definir in "A définir":
            feat_pool, feat = "A définir" # Extraire features pour chaque batch

        gall_feat_pool[ptr:ptr + batch_size, :] = feat_pool.detach().cpu().numpy()
        gall_feat[ptr:ptr + batch_size, :] = feat.detach().cpu().numpy()
        ptr = ptr + batch_size

    return gall_feat_pool, gall_feat

if __name__ == "__main__":

    # Work on GPU if available
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    # Import the lfw dataset - With 15 img min per person
    print("===> Loading Fetch Dataset \n")
    lfw_people = fetch_lfw_people(min_faces_per_person=15, resize=0.4)
    nb_img_per_id_to_keep = 15

    # Renvoie la base de données avec 15 images par personnes maximum
    # Vous devrez coder la fonction dans le fichier utils
    X, y, n_classes = prepare_dataset(lfw_people, nb_img_per_id_to_keep)

    # Produire les listes d'identités qui correspondront aux folds d'entraînement / validation et
    # aux identités de test.
    # train_ids_lists et val_ids_lists seront des listes de 5 listes d'identités (Une liste pour chaque fold)
    # test_ids sera une liste de 21 identités
    train_ids_lists, val_ids_lists, test_ids = prepare_data_ids(n_classes)

    # Prepare var - Batch size - On utilise 4 images par identités et 8 identités par batch (=> 32 imgs)
    num_img_of_same_id_in_batch = 4
    num_different_identities_in_batch = 8
    batch_size = num_img_of_same_id_in_batch * num_different_identities_in_batch # batch_size = 32
    test_batch_size = 32

    global_img_pos = GenIdx(y)  # Get the images positions in list for each specific identity

    folds = "A définir"  # Number of fold K for cross-validation
    epochs = "A définir" # Number of epochs

    # Augmentation de la taille des images
    img_h, img_w = 288, 144

    # Préparation des transformations appliquées aux images pour l'entraînement ou le reste (validation / test)
    transform_train = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((img_h, img_w)),
        transforms.Pad(10),
        transforms.RandomCrop((img_h, img_w)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
    ])

    transform_test = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((img_h, img_w)),
        transforms.ToTensor(),
    ])

    # Définir votre ou vos fonctions de pertes
    # Vous pouvez définir votre/vos fonction de pertes dans le fichier loss.py
    criterion_1 = "A définir"

    for fold in range(folds):
        net = model(n_classes).to(device)

        # Définir votre optimizer :
        opt = "A définir"

        # Préparation du query / gallery set pour la validation
        # A noter que query set = gallery set si on compte comparer chaque image de la base avec toute les autres
        query_gall_set, n_query_gall = prepare_set(X, global_img_pos, val_ids_lists[fold], transform=transform_test)
        query_gall_loader = torch.utils.data.DataLoader(query_gall_set, batch_size=test_batch_size, shuffle=False)

        best_map = 0
        for epoch in range(epochs):

            trainset = LFW_training_Data(X, global_img_pos, train_ids_lists[fold], transform=transform_train)

            img_pos = GenIdx(trainset.train_labels)  # Get the images positions in list for each specific identity

            # Permet de sampler la base de données par batch contenant le bon nombre d'identités et d'images par identités
            sampler = IdentitySampler(trainset.train_labels, img_pos, num_img_of_same_id_in_batch,
                                      num_different_identities_in_batch)

            trainset.cIndex = sampler.index

            trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, \
                                                      sampler=sampler, drop_last=True)

            train(epoch, criterion_1, opt, trainloader, device, net)

            # Call the validation every two epochs
            if epoch != 0 and epoch % 2 == 0:
                cmc, mAP, cmc_att, mAP_att = valid(query_gall_loader, query_gall_set.val_test_label)

                # Save model based on validation mAP
                if mAP > best_map:  # Usual saving

                    best_map = mAP
                    best_epoch = epoch
                    state = {
                        'net': net.state_dict(),
                        'cmc': cmc,
                        'mAP': mAP,
                        'epoch': epoch,
                    }
                    torch.save(state, f'model_fold({fold})_best.t')

                print('fc : Rank-1: {:.2%} | Rank-5: {:.2%} | mAP: {:.2%}'.format(cmc[0], cmc[4], mAP ))
                print('att : Rank-1: {:.2%} | Rank-5: {:.2%} | mAP: {:.2%}'.format(cmc_att[0], cmc_att[4], mAP_att))

    # Tester ici ou dans un fichier à part votre modèle sur les données de test.


import collections
import csv
import os
from metadata import *
from glob import glob
import json
from fuzzywuzzy import fuzz
import textdistance
import spacy
nlp = spacy.load("en_core_web_sm")

class InputExample(object):
  """A single training/test example for simple sequence classification."""

  def __init__(self, guid, text_a, text_b=None, label=None):
    """Constructs a InputExample.

    Args:
      guid: Unique id for the example.
      text_a: string. The untokenized text of the first sequence. For single
        sequence tasks, only this sequence must be specified.
      text_b: (Optional) string. The untokenized text of the second sequence.
        Only must be specified for sequence pair tasks.
      label: (Optional) string. The label of the example. This should be
        specified for train and dev examples, but not for test examples.
    """
    self.guid = guid
    self.text_a = text_a
    self.text_b = text_b
    self.label = label

class DataProcessor(object):
  """Base class for data converters for sequence classification data sets."""

  def get_train_examples(self, data_dir):
    """Gets a collection of `InputExample`s for the train set."""
    raise NotImplementedError()

  def get_dev_examples(self, data_dir):
    """Gets a collection of `InputExample`s for the dev set."""
    raise NotImplementedError()

  def get_test_examples(self, data_dir):
    """Gets a collection of `InputExample`s for prediction."""
    raise NotImplementedError()

  def get_labels(self):
    """Gets the list of labels for this data set."""
    raise NotImplementedError()

  @classmethod
  def _read_tsv(cls, input_file, quotechar=None):
    """Reads a tab separated value file."""
    with open(input_file, "r", encoding="utf-8") as f:
      reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
      lines = []
      for line in reader:
        lines.append(line)
      return lines

class SciProcessor(DataProcessor):
    def get_train_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

    def get_dev_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")

    def get_test_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")

    def get_labels(self):
        """See base class."""
        return ["true", "false"]

    def tokenizer(self,sentence):
        return sentence.split()

    def get_test_result(self, testdata_dir, output_dir):
        testlines = self._read_tsv(os.path.join(testdata_dir, "test.tsv"))
        testresult = self._read_tsv(os.path.join(output_dir, "test_results.tsv"))
        prediction_mention = {}
        prediction_value = {}
        prediction_ground = {}
        for (i, line) in enumerate(testlines):
            line = line + testresult[i - 1]
            if line[1] in prediction_mention:
                if (float(line[6]) > prediction_value.get(line[1])):
                    prediction_mention[line[1]] = line[2]
                    prediction_value[line[1]] = float(line[6])
                    prediction_ground[line[1]] = line[0]
            else:
                prediction_mention[line[1]] = line[2]
                prediction_value[line[1]] = float(line[6])
                prediction_ground[line[1]] = line[0]
        for i, j in prediction_mention.items():
            print(i, j)
        tp = 0
        for i in prediction_ground.values():
            if i == 1:
                tp = tp + 1
        print(prediction_ground.keys().__len__(), tp)
        print("evaluation finished")

    def _create_examples(self, lines, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        for (i, line) in enumerate(lines):
            guid = "%s-%s" % (set_type, i)
            text_a = self.tokenizer(line[2])
            text_b = self.tokenizer(line[3])
            if set_type == "test":
                label = "true"
            else:
                label = self.tokenizer(line[0])
            examples.append(
                InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
        return examples

class KBReader(object):
    def __init__(self):
        self.dataset_dict = self.read_datasetKB()
        self.metric_dict = self.read_metricKB()
        self.task_dict = self.read_taskKB()

    def read_datasetKB(self):
        dataset_dict = dict()
        with open(datasetKB_path) as f:
            reader = csv.reader(f, delimiter="\t")
            for line in reader:
                dataset_dict[line[0]] = line[1]
        return dataset_dict

    def read_metricKB(self):
        metric_dict = dict()
        with open(metricKB_path) as f:
            reader = csv.reader(f, delimiter="\t")
            for line in reader:
                metric_dict[line[0]] = line[1]
        return metric_dict

    def read_taskKB(self):
        task_dict = dict()
        with open(taskKB_path) as f:
            reader = csv.reader(f, delimiter="\t")
            for line in reader:
                task_dict[line[0]] = line[1]
        return task_dict

class PaperReader(object):
    def __init__(self,corpus='NLP-TDMS'):
        self.corpus = corpus
        self.paper_text,self.paper_table = self.read_papers()

    def get_paperIDs(self):
        paper_fnames = []
        with open(os.path.join(NLP_TDM_path,'downloader','paper_links.tsv')) as f:
            reader = csv.reader(f, delimiter="\t")
            for line in reader:
                paper_fnames.append(line[0])
        return paper_fnames

    def get_PaperTxt_files(self):
        return glob(os.path.join(NLP_TDM_path,'pdfFile_txt','*'))

    def get_PaperTable_files(self):
        return glob(os.path.join(NLP_TDM_path, 'pdfFile_table', '*'))

    def read_papers(self):
        pids = self.get_paperIDs()
        text_fnames = self.get_PaperTxt_files()
        text_dict = dict()
        for t_fname in text_fnames:
            text_dict[os.path.basename(t_fname).split('.txt')[0]] = t_fname

        table_fnames = self.get_PaperTable_files()
        table_dict = dict()
        for t_fname in table_fnames:
            table_dict[os.path.basename(t_fname).split('.json')[0]] = t_fname

        paper_text = dict()
        paper_table = dict()
        for pid in pids:
            pid = pid.split('.pdf')[0]
            if pid in text_dict:
                f = open(text_dict[pid],'r')
                paper_text[pid] = f.read()
                f.close()
            if pid in table_dict:
                f = open(table_dict[pid],'r')
                paper_table[pid] = json.loads(f.readline())
                f.close()
        return paper_text,paper_table

    def get_paper_datasets(self):
        paper_dataset = dict()
        with open(os.path.join(NLP_TDM_path, 'annotations', 'datasetAnnotation.tsv')) as f:
            reader = csv.reader(f, delimiter="\t")
            for line in reader:
                paper_dataset[line[0].split('.pdf')[0]] = line[1]
        return paper_dataset

    def get_paper_task(self):
        paper_task = dict()
        with open(os.path.join(NLP_TDM_path, 'annotations', 'taskAnnotation.tsv')) as f:
            reader = csv.reader(f, delimiter="\t")
            for line in reader:
                paper_task[line[0].split('.pdf')[0]] = line[1]
        return paper_task

    def get_paper_result(self):
        #each paper entry has the format:
        #task1#dataset1#metric1#score1$task2#dataset2#metric2#score2$...
        paper_result = dict()
        with open(os.path.join(NLP_TDM_path, 'annotations', 'resultsAnnotation.tsv')) as f:
            reader = csv.reader(f, delimiter="\t")
            for line in reader:
                paper_result[line[0].split('.pdf')[0]] = line[1]
        return paper_result

    def get_num_pos(self):
        ct = 0
        paper_result = self.get_paper_result()
        for pid in paper_result:
            results = paper_result[pid]
            results = results.split('$')
            ct += len(results)
        return ct

def createTagData():
    '''
    for each paper id:
        get paper result
        get task, dataset, metric, score from result
        matching those entities in text and record the token
        scores should be matched only in tables
    '''
    sim_thred = 50
    data_thred = 100
    output_mode = 'key_sent' # all_sec / key_sec
    preader = PaperReader()
    paper_result = preader.get_paper_result()
    paper_text, paper_table = preader.read_papers()
    f_out = open('./data/all.data','w')
    for pid in paper_result:
        print(pid)
        if pid not in paper_text:
            continue
        text = paper_text[pid]
        lines = text.split('\n')
        data_pairs = []
        for line in lines:
            if len(line) == 0:
                continue
            doc = nlp(line)
            # first check Sections
            if doc[0].text == 'section' and doc[1].text == ':':
                for sec_idx in range(2,len(doc)):
                    if sec_idx == 2:
                        if len(doc) > 3:
                            data_pairs.append([doc[sec_idx].text,'SECTITLE_START'])
                        else:
                            data_pairs.append([doc[sec_idx].text, 'SECTITLE_END'])
                    elif sec_idx + 1< len(doc):
                        data_pairs.append([doc[sec_idx].text, 'SECTITLE_CONTENT'])
                    elif sec_idx + 1 == len(doc):
                        data_pairs.append([doc[sec_idx].text, 'SECTITLE_END'])
                continue

            entity_entries = []  # (start,end,entity type)
            results = paper_result[pid]
            results = results.split('$')
            has_dataset = False
            for result in results:
                seps = result.split('#')
                task,dataset,metric,score = seps[:4]
                #print(seps)
                # fuzzy matching for entities
                for np in doc.noun_chunks:
                    sim = fuzz.ratio(np.text.lower(),task.replace('_',' ').lower())
                    #sim = textdistance.jaccard(np.text.lower(),task.replace('_',' ').lower())
                    if sim > sim_thred:
                        entity_entries.append([np.start,np.end,'task',task])
                    sim = fuzz.ratio(np.text.lower(), dataset.lower())
                    #sim = textdistance.jaccard(np.text.lower(), dataset.replace('_',' ').lower())
                    if sim == data_thred:
                        entity_entries.append([np.start, np.end, 'dataset',dataset])
                        has_dataset= True
                        #print("{0}:{1}:{2}".format(np.text,dataset,sim ))
                    sim = fuzz.ratio(np.text.lower(), metric.replace('_',' ').lower())
                    #sim = textdistance.jaccard(np.text.lower(), metric.replace('_',' ').lower())
                    if sim > sim_thred:
                        entity_entries.append([np.start, np.end, 'metric',metric])
                        #print("{0}:{1}:{2}".format(np.text, metric, sim))

            entity_entries.sort(key=lambda x: x[0])

            #begin to output every token
            if len(entity_entries) == 0:
                current_entity = [len(doc)+1,len(doc)+1]
            else:
                current_entity = entity_entries.pop(0)

            if output_mode == 'all_sec':
                idx = 0
                while idx < len(doc):
                    if idx == 0:
                        data_pairs.append([doc[idx].text,'SEC_START'])
                    elif idx + 1 < len(doc):
                        #current token belongs to an entity
                        if current_entity[0] == idx:
                            data_pairs.append([doc[idx].text, current_entity[2]+"START"])
                        elif idx ==  current_entity[1]-1:
                            data_pairs.append([doc[idx].text, current_entity[2] + "END"])
                        elif  current_entity[0] < idx+1 <  current_entity[1]:
                            data_pairs.append([doc[idx].text,current_entity[2]])
                        else:
                            data_pairs.append([doc[idx].text, 'SEC_CONTENT'])
                        #finish tag the entity
                        if idx+1 == current_entity[1] and len(entity_entries) != 0:
                            current_entity = entity_entries.pop(0)

                    elif idx + 1 == len(doc):
                        data_pairs.append([doc[idx].text, 'SEC_END'])
                    idx += 1
            if output_mode == 'key_sent':
                sents = list(doc.sents)
                for sent in sents:
                    has_ent = False
                    if len(sent) < 3:
                        continue
                    #if the sentence does not have an entity, skip
                    if sent.start <= current_entity[0] and sent.end >=  current_entity[1]:
                        idx = sent.start

                        while idx < sent.end:
                            if idx == sent.start and idx != current_entity[0]:
                                data_pairs.append([sent[idx-sent.start].text, 'SENT_START'])
                            elif idx + 1 < sent.end:
                                # current token belongs to an entity
                                if current_entity[0] == idx:
                                    has_ent = True
                                    #data_pairs.append([sent[idx-sent.start].text, current_entity[2] + "START"])
                                    if current_entity[2] == 'task':
                                        data_pairs.append([current_entity[3], current_entity[2]])
                                    else:
                                        data_pairs.append([sent[idx - sent.start].text, current_entity[2]])
                                    print("{0}:{1}".format(current_entity[2],current_entity[3]))
                                elif idx == current_entity[1] - 1:
                                    if current_entity[2] != 'task':
                                    #data_pairs.append([sent[idx-sent.start].text, current_entity[2] + "END"])
                                        data_pairs.append([sent[idx - sent.start].text, current_entity[2]])
                                elif current_entity[0] < idx + 1 < current_entity[1]:
                                    if current_entity[2] != 'task':
                                    #data_pairs.append([sent[idx-sent.start].text, current_entity[2]])
                                        data_pairs.append([sent[idx-sent.start].text, current_entity[2]])
                                else:
                                    data_pairs.append([sent[idx-sent.start].text, 'SENT_CONTENT'])
                                # finish tag the entitys
                                if idx + 1 == current_entity[1] and len(entity_entries) != 0:
                                    current_entity = entity_entries.pop(0)

                            elif idx + 1 ==  sent.end and idx+1 != current_entity[1]:
                                data_pairs.append([sent[idx-sent.start].text, 'SENT_END'])
                            idx += 1

        #write into files
        f_out.write(pid + '\t')
        for pair in data_pairs:
            f_out.write(pair[0] + '\\tag' + pair[1] + '\t')
            #print(pair[0] + '\t' + pair[1])
        f_out.write('\n')

    f_out.close()

if __name__ == '__main__':
    createTagData()


from branca.element import MacroElement
from jinja2 import Template

from folium.elements import JSCSSMixin
from folium.utilities import parse_options

class MousePosition(JSCSSMixin, MacroElement):
    """Add a field that shows the coordinates of the mouse position.

    Uses the Leaflet plugin by Ardhi Lukianto under MIT license.
    https://github.com/ardhi/Leaflet.MousePosition

    Parameters
    ----------
    position : str, default 'bottomright'
        The standard Control position parameter for the widget.
    separator : str, default ' : '
        Character used to separate latitude and longitude values.
    empty_string : str, default 'Unavailable'
       Initial text to display.
    lng_first : bool, default False
        Whether to put the longitude first or not.
        Set as True to display longitude before latitude.
    num_digits : int, default '5'
        Number of decimal places included in the displayed
        longitude and latitude decimal degree values.
    prefix : str, default ''
        A string to be prepended to the coordinates.
    lat_formatter : str, default None
        Custom Javascript function to format the latitude value.
    lng_formatter : str, default None
        Custom Javascript function to format the longitude value.

    Examples
    --------
    >>> fmtr = "function(num) {return L.Util.formatNum(num, 3) + ' º ';};"
    >>> MousePosition(
    ...     position="topright",
    ...     separator=" | ",
    ...     prefix="Mouse:",
    ...     lat_formatter=fmtr,
    ...     lng_formatter=fmtr,
    ... )

    """

    _template = Template(
        """
        {% macro script(this, kwargs) %}
            var {{ this.get_name() }} = new L.Control.MousePosition(
                {{ this.options|tojson }}
            );
            {{ this.get_name() }}.options["latFormatter"] =
                {{ this.lat_formatter }};
            {{ this.get_name() }}.options["lngFormatter"] =
                {{ this.lng_formatter }};
            {{ this._parent.get_name() }}.addControl({{ this.get_name() }});
        {% endmacro %}
    """
    )

    default_js = [
        (
            "Control_MousePosition_js",
            "https://cdn.jsdelivr.net/gh/ardhi/Leaflet.MousePosition/src/L.Control.MousePosition.min.js",
        )
    ]
    default_css = [
        (
            "Control_MousePosition_css",
            "https://cdn.jsdelivr.net/gh/ardhi/Leaflet.MousePosition/src/L.Control.MousePosition.min.css",
        )
    ]

    def __init__(
        self,
        position="bottomright",
        separator=" : ",
        empty_string="Unavailable",
        lng_first=False,
        num_digits=5,
        prefix="",
        lat_formatter=None,
        lng_formatter=None,
        **kwargs
    ):
        super().__init__()
        self._name = "MousePosition"

        self.options = parse_options(
            position=position,
            separator=separator,
            empty_string=empty_string,
            lng_first=lng_first,
            num_digits=num_digits,
            prefix=prefix,
            **kwargs
        )
        self.lat_formatter = lat_formatter or "undefined"
        self.lng_formatter = lng_formatter or "undefined"


from django.http import HttpRequest
from django.contrib.auth import get_user_model
from django.utils.translation import gettext_lazy as _

from ninja import Router, Query

from apps.base.schemas import FormInvalidResponseContract
from skii.endpoint.filters import LessonFilterSchema

from skii.platform.models.agent import TeacherAgent, StudentAgent
from skii.endpoint.schemas.response import TeacherLessonContract
from skii.endpoint.schemas.response import StudentLessonContract
from skii.endpoint.schemas.identifier import IntStrUUID4
from skii.platform.models.event import LessonEvent

UserModel = get_user_model()
MsgErrorStudent = _(
    "Fetch student lessons with another student account "
    "logged is forbidden. Please contact your support."
)

route_agenda = Router(tags=["skii", "agenda"])

RouterModel = LessonEvent

@route_agenda.get(
    path="/teacher_lessons/{teacher_pk}/",
    response={
        200: TeacherLessonContract,
        422: FormInvalidResponseContract,
    },
)
def teacher_lessons(
    request: HttpRequest,
    teacher_pk: IntStrUUID4,
    filters: LessonFilterSchema = Query(...),
):
    """Fetch teacher with them lessons assigned.

    This view permit's to use query filters:
        - start: Datetime to filter lesson start up to
        - stop: Datetime to filter lesson stop down to
    """
    agent = TeacherAgent.objects.get(pk=teacher_pk)
    lessons = filters.filter(LessonEvent.objects.filter(teacher__pk=agent.pk))
    setattr(agent, "lessons_assigned", lessons)
    return 200, agent

@route_agenda.get(
    path="/student_lessons/{student_pk}/",
    response={
        200: StudentLessonContract,
        422: FormInvalidResponseContract,
    },
)
def student_lessons(
    request: HttpRequest,
    student_pk: IntStrUUID4,
    filters: LessonFilterSchema = Query(...),
):
    """Fetch student with them lessons reserved.

    This view permit's to use query filters:
        - start: Datetime to filter lesson start up to
        - stop: Datetime to filter lesson stop down to

    They can be used independently, or together.

    Security:   Pending an implementation of more advanced rights,
                we will simply prohibit lesson requests from a student
                other than the connected user.
    """
    agent = StudentAgent.objects.get(pk=student_pk)
    # If the logged user is a Student (filter(user=request.user).exist())
    # and different than the requested one (request.user)
    # Raise a Permission error.
    if (
        request.user != agent.user
        and StudentAgent.objects.filter(user=request.user).exists()
    ):
        raise PermissionError(MsgErrorStudent)

    lessons = filters.filter(LessonEvent.objects.filter(students__pk=agent.pk))
    setattr(agent, "lessons_subscribed", lessons)

    return 200, agent


from ariadne import convert_kwargs_to_snake_case

def shopping_list_items_resolver(obj, info):
    from api.models import ShoppingListItem
    try:
        items = [item.to_dict() for item in ShoppingListItem.query.all()]
        print(items)
        payload = {
            "success": True,
            "shopping_list_items": items
        }
    except Exception as error:
        payload = {
            "success": False,
            "errors": [str(error)]
        }
    return payload

@convert_kwargs_to_snake_case
def get_shopping_list_item_resolver(obj, info, uid):
    from api.models import ShoppingListItem
    try:
        item = ShoppingListItem.query.get(uid)
        print(item)
        payload = {
            "success": True,
            "shopping_list_item": item.to_dict()
        }
    except AttributeError as error:
        payload = {
            "success": False,
            "errors": [f"Не найден uid {uid}"]
        }
    except Exception as error:
        payload = {
            "success": False,
            "errors": [str(error)]
        }
    return payload


import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim

from networks.network_td3 import Actor, Critic
from common.buffer import Replay_buffer as buffer
from pathlib import Path
import sys
import os

base_dir = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(base_dir))

max_action = 2.0
min_Val = torch.tensor(1e-7).float()  # min value

def get_trajectory_property():
    return ["action", "logits"]

class TD3:
    def __init__(self, args):

        self.state_dim = args.obs_space
        self.action_dim = 1

        self.actor = Actor(self.state_dim, self.action_dim, max_action)
        self.actor_target = Actor(self.state_dim, self.action_dim, max_action)
        self.critic_1 = Critic(self.state_dim, self.action_dim)
        self.critic_1_target = Critic(self.state_dim, self.action_dim)
        self.critic_2 = Critic(self.state_dim, self.action_dim)
        self.critic_2_target = Critic(self.state_dim, self.action_dim)

        self.actor_optimizer = optim.Adam(self.actor.parameters())
        self.critic_1_optimizer = optim.Adam(self.critic_1.parameters())
        self.critic_2_optimizer = optim.Adam(self.critic_2.parameters())

        self.actor_target.load_state_dict(self.actor.state_dict())
        self.critic_1_target.load_state_dict(self.critic_1.state_dict())
        self.critic_2_target.load_state_dict(self.critic_2.state_dict())

        self.max_action = max_action
        self.buffer_size = args.buffer_capacity
        self.batch_size = args.batch_size
        self.gamma = args.gamma
        self.tau = args.tau
        self.update_freq = args.update_freq
        self.eps = 0.2

        self.policy_noise = args.policy_noise
        self.noise_clip = args.noise_clip
        self.policy_delay = args.policy_delay
        self.exploration_noise = args.exploration_noise

        trajectory_property = get_trajectory_property()
        self.memory = buffer(self.buffer_size, trajectory_property)
        self.memory.init_item_buffers()

        self.num_critic_update_iteration = 0
        self.num_actor_update_iteration = 0
        self.num_training = 0

    def choose_action(self, observation, train=True):
        inference_output = self.inference(observation, train)
        if train:
            self.add_experience(inference_output)
        return inference_output

    def inference(self, state, train=True):
        state = torch.tensor(state, dtype=torch.float).view(1, -1)
        action = self.actor(state).cpu().data.numpy().flatten()
        action += np.random.normal(0, self.exploration_noise, size=self.action_dim)
        action = action.clip(-max_action, max_action)
        logits = self.actor(state).detach().numpy()
        return {"action": action, "logits": logits}

    def add_experience(self, output):
        agent_id = 0
        for k, v in output.items():
            self.memory.insert(k, agent_id, v)

    def learn(self):

        data_length = len(self.memory.item_buffers["rewards"].data)
        if data_length < self.buffer_size:
            return

        for i in range(self.update_freq):
            data = self.memory.sample(self.batch_size)
            transitions = {
                "o_0": np.array(data["states"]),
                "o_next_0": np.array(data["states_next"]),
                "r_0": np.array(data["rewards"]).reshape(-1, 1),
                "u_0": np.array(data["action"]),
                "d_0": np.array(data["dones"]).reshape(-1, 1),
            }

            state = torch.tensor(transitions["o_0"], dtype=torch.float)
            next_state = torch.tensor(transitions["o_next_0"], dtype=torch.float)
            action = torch.tensor(transitions["u_0"], dtype=torch.float).view(
                self.batch_size, -1
            )
            reward = torch.tensor(transitions["r_0"], dtype=torch.float).view(
                self.batch_size, -1
            )
            done = torch.tensor(transitions["d_0"], dtype=torch.float).view(
                self.batch_size, -1
            )
            # Select next action according to target policy:

            noise = torch.ones_like(action).data
            noise = torch.tensor(noise, dtype=torch.float)
            noise = noise.normal_(0.0, self.policy_noise)
            noise = noise.clamp(-self.noise_clip, self.noise_clip)
            next_action = self.actor_target(next_state) + noise
            next_action = next_action.clamp(-self.max_action, self.max_action)

            # Compute target Q-value:
            target_Q1 = self.critic_1_target(next_state, next_action)
            target_Q2 = self.critic_2_target(next_state, next_action)
            target_Q = torch.min(target_Q1, target_Q2)
            target_Q = reward + ((1 - done) * self.gamma * target_Q).detach()

            # Optimize Critic 1:
            current_Q1 = self.critic_1(state, action)
            loss_Q1 = F.mse_loss(current_Q1, target_Q)
            self.critic_1_optimizer.zero_grad()
            loss_Q1.backward()
            self.critic_1_optimizer.step()

            # Optimize Critic 2:
            current_Q2 = self.critic_2(state, action)
            loss_Q2 = F.mse_loss(current_Q2, target_Q)
            self.critic_2_optimizer.zero_grad()
            loss_Q2.backward()
            self.critic_2_optimizer.step()

            # Delayed policy updates:
            if i % self.policy_delay == 0:
                # Compute actor loss:
                actor_loss = -self.critic_1(state, self.actor(state)).mean()

                # Optimize the actor
                self.actor_optimizer.zero_grad()
                actor_loss.backward()
                self.actor_optimizer.step()
                for param, target_param in zip(
                    self.actor.parameters(), self.actor_target.parameters()
                ):
                    target_param.data.copy_(
                        ((1 - self.tau) * target_param.data) + self.tau * param.data
                    )

                for param, target_param in zip(
                    self.critic_1.parameters(), self.critic_1_target.parameters()
                ):
                    target_param.data.copy_(
                        ((1 - self.tau) * target_param.data) + self.tau * param.data
                    )

                for param, target_param in zip(
                    self.critic_2.parameters(), self.critic_2_target.parameters()
                ):
                    target_param.data.copy_(
                        ((1 - self.tau) * target_param.data) + self.tau * param.data
                    )

                self.num_actor_update_iteration += 1
        self.num_critic_update_iteration += 1
        self.num_training += 1

    def save(self, save_path, episode):
        base_path = os.path.join(save_path, "trained_model")
        if not os.path.exists(base_path):
            os.makedirs(base_path)

        model_critic_1_path = os.path.join(
            base_path, "critic_1" + str(episode) + ".pth"
        )
        torch.save(self.critic_1.state_dict(), model_critic_1_path)
        model_critic_1_target_path = os.path.join(
            base_path, "critic_1_target" + str(episode) + ".pth"
        )
        torch.save(self.critic_1_target.state_dict(), model_critic_1_target_path)

        model_critic_2_path = os.path.join(
            base_path, "critic_2" + str(episode) + ".pth"
        )
        torch.save(self.critic_2.state_dict(), model_critic_2_path)
        model_critic_2_target_path = os.path.join(
            base_path, "critic_2_target" + str(episode) + ".pth"
        )
        torch.save(self.critic_2_target.state_dict(), model_critic_2_target_path)

        model_actor_path = os.path.join(base_path, "actor" + str(episode) + ".pth")
        torch.save(self.actor.state_dict(), model_actor_path)
        model_actor_target_path = os.path.join(
            base_path, "actor_target" + str(episode) + ".pth"
        )
        torch.save(self.actor_target.state_dict(), model_actor_target_path)

    def load(self, save_path):
        self.actor.load_state_dict(torch.load(str(save_path) + "actor.pth"))
        self.actor_target.load_state_dict(
            torch.load(str(save_path) + "actor_target.pth")
        )
        self.critic_1.load_state_dict(torch.load(str(save_path) + "critic_1.pth"))
        self.critic_1_target.load_state_dict(
            torch.load(str(save_path) + "critic_1_target.pth")
        )
        self.critic_2.load_state_dict(torch.load(str(save_path) + "critic_2.pth"))
        self.critic_2_target.load_state_dict(
            torch.load(str(save_path) + "critic_2_target.pth")
        )


# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import json
import subprocess

cached_functions = {}

def getMetricsJson(src_url):
    if src_url.startswith("http"):
        print("Fetching source for function extraction: {}".format(src_url))
        metrics = subprocess.check_output(["./fetch_fn_names.sh", src_url])
    else:
        print("Skip fetching source: {}".format(src_url))
        metrics = ""

    try:
        return json.loads(metrics)
    except ValueError:
        return {"kind": "empty", "name": "anonymous", "spaces": []}

def getSpaceFunctionsRecursive(metrics_space):
    functions = []
    if (
        metrics_space["kind"] == "function"
        and metrics_space["name"]
        and metrics_space["name"] != "<anonymous>"
    ):
        functions.append(
            {
                "name": metrics_space["name"],
                "start_line": int(metrics_space["start_line"]),
                "end_line": int(metrics_space["end_line"]),
            }
        )
    for space in metrics_space["spaces"]:
        functions += getSpaceFunctionsRecursive(space)
    return functions

def getSourceFunctions(src_url):
    if src_url not in cached_functions:
        metrics_space = getMetricsJson(src_url)
        cached_functions[src_url] = getSpaceFunctionsRecursive(metrics_space)

    return cached_functions[src_url]

def getFunctionName(location):
    location.replace("annotate", "raw-file")
    pieces = location.split("#l")
    src_url = pieces[0]
    line = int(pieces[1])
    closest_name = "<Unknown {}>".format(line)
    closest_start = 0
    functions = getSourceFunctions(src_url)
    for fn in functions:
        if (
            fn["start_line"] > closest_start
            and line >= fn["start_line"]
            and line <= fn["end_line"]
        ):
            closest_start = fn["start_line"]
            closest_name = fn["name"]
    return closest_name


from pathlib import Path

from utils.io_lib import read_file_to_string, puzzle_input_file_name, parse_input_string_to_groups

DAY_NUMBER = 1

def task1():
    """
    Find the Elf carrying the most Calories. How many total Calories is that Elf carrying?
    """
    print(f"Maximum Calories carried by an Elf is {max(calculate_sums_of_lists())}")

def calculate_sums_of_lists() -> list[int]:
    file_input_string = read_file_to_string(puzzle_input_file_name(day_number=DAY_NUMBER))
    list_of_strings = parse_input_string_to_groups(file_input_string)
    sums_of_lists = []
    for index, value in enumerate(list_of_strings):
        value_list = value.splitlines()
        sums_of_lists.append(sum(list(map(int, value_list))))
    return sums_of_lists

def task2():
    """
    Find the top three Elves carrying the most Calories. How many Calories are those Elves carrying in total?
    """
    print(f"Sum of Calories carried by top three Elves is {sum(sorted(calculate_sums_of_lists(), reverse=True)[:3])}")

if __name__ == "__main__":
    task1()
    task2()