import os
import stat
import inspect
import tempfile
import copy
import pyfastaq
import shutil
import multiprocessing
import collections
import iva

class Error (Exception): pass

class Qc:
    def __init__(self,
        assembly_fasta,
        output_prefix,
        embl_dir=None,
        ref_db=None,
        reads_fr=None,
        reads_fwd=None,
        reads_rev=None,
        ratt_config=None,
        min_ref_cov=5,
        contig_layout_plot_title="IVA QC contig layout and read depth",
        threads=1,
        nucmer_min_cds_hit_length=20,
        nucmer_min_cds_hit_id=80,
        nucmer_min_ctg_hit_length=100,
        nucmer_min_ctg_hit_id=80,
        gage_nucmer_minid=80,
        smalt_k=15,
        smalt_s=3,
        smalt_id=0.5,
        blast_for_act=False,
        kraken_preload=False,
        clean=True,
    ):

        if embl_dir is None and ref_db is None:
           raise Error('Must provide either embl_dir or ref_db to Qc object. Cannot continue')

        self.embl_dir = embl_dir
        if self.embl_dir is not None:
            self.embl_dir = os.path.abspath(self.embl_dir)

        self.embl_files = []

        self.ref_db = ref_db

        files_to_check = [assembly_fasta, reads_fr, reads_fwd, reads_rev]
        for filename in files_to_check:
            if filename is not None and not os.path.exists(filename):
                raise Error('Error in IVA QC. File not found: "' + filename + '"')

        self.outprefix = output_prefix
        self.assembly_fasta = output_prefix + '.assembly.fasta'
        self.assembly_bam = output_prefix + '.reads_mapped_to_assembly.bam'
        self.ref_bam = output_prefix + '.reads_mapped_to_ref.bam'
        self.reads_fwd = reads_fwd
        self.reads_rev = reads_rev
        self.threads = threads
        self.kraken_preload = kraken_preload
        self.kraken_prefix = self.outprefix + '.kraken'
        self.ratt_config = None if ratt_config is None else os.path.abspath(ratt_config)
        self.ratt_outdir = self.outprefix + '.ratt'
        self.blast_for_act = blast_for_act
        self.blast_out = self.outprefix + '.assembly_v_ref.blastn'
        self.fasta_assembly_contigs_hit_ref = self.outprefix + '.assembly_contigs_hit_ref.fasta'
        self.fasta_assembly_contigs_not_hit_ref = self.outprefix + '.assembly_contigs_not_hit_ref.fasta'
        self.act_script = self.outprefix + '.assembly_v_ref.act.sh'
        self.gage_outdir = self.outprefix + '.gage'
        self.gage_nucmer_minid = gage_nucmer_minid
        self.files_to_clean = []
        self.clean = clean

        if reads_fr:
            self.reads_fwd = self.outprefix + '.reads_1'
            self.reads_rev = self.outprefix + '.reads_2'
            self.files_to_clean.append(self.reads_fwd)
            self.files_to_clean.append(self.reads_rev)
            pyfastaq.tasks.deinterleave(reads_fr, self.reads_fwd, self.reads_rev)

        if not (None not in [self.reads_fwd, self.reads_rev] or reads_fr is not None):
            raise Error('IVA QC needs reads_fr or both reads_fwd and reads_rev')

        self._set_assembly_fasta_data(assembly_fasta)

        def unzip_file(infile, outfile):
            iva.common.syscall('gunzip -c ' + infile + ' > ' + outfile)

        processes = []

        if self.reads_fwd.endswith('.gz'):
            new_reads_fwd = self.outprefix + '.reads_1'
            processes.append(multiprocessing.Process(target=unzip_file, args=(self.reads_fwd, new_reads_fwd)))
            self.reads_fwd = new_reads_fwd
            self.files_to_clean.append(self.reads_fwd)

        if self.reads_rev.endswith('.gz'):
            new_reads_rev = self.outprefix + '.reads_2'
            processes.append(multiprocessing.Process(target=unzip_file, args=(self.reads_rev, new_reads_rev)))
            self.reads_rev = new_reads_rev
            self.files_to_clean.append(self.reads_rev)

        if len(processes) == 1:
            processes[0].start()
            processes[0].join()
        elif len(processes) == 2:
            processes[0].start()
            if self.threads == 1:
                processes[0].join()
            processes[1].start()
            processes[1].join()
            if self.threads > 1:
                processes[0].join()

        self.min_ref_cov = min_ref_cov
        self.threads = threads
        self.contig_layout_plot_title = contig_layout_plot_title
        self.nucmer_min_cds_hit_length = nucmer_min_cds_hit_length
        self.nucmer_min_cds_hit_id = nucmer_min_cds_hit_id
        self.nucmer_min_ctg_hit_length = nucmer_min_ctg_hit_length
        self.nucmer_min_ctg_hit_id = nucmer_min_ctg_hit_id
        self.smalt_k=smalt_k
        self.smalt_s=smalt_s
        self.smalt_id=smalt_id
        self.contig_pos_in_ref = {}
        self.low_cov_ref_regions = {}
        self.low_cov_ref_regions_fwd = {}
        self.low_cov_ref_regions_rev = {}
        self.ok_cov_ref_regions = {}
        self.ref_coverage_fwd = {}
        self.ref_coverage_rev = {}
        self.ref_cds_fasta = output_prefix + '.ref_cds_seqs.fa'
        self.cds_nucmer_coords_in_assembly = output_prefix + '.ref_cds_seqs_mapped_to_assembly.coords'
        self.cds_assembly_stats = {}
        self.refseq_assembly_stats = {}
        self.assembly_vs_ref_coords = output_prefix + '.assembly_vs_ref.coords'
        self.assembly_vs_ref_mummer_hits = {}
        self.ref_pos_covered_by_contigs = {}
        self.ref_pos_not_covered_by_contigs = {}
        self.should_have_assembled = {}
        self.contig_placement = {}
        self.stats_keys = [
            'ref_EMBL_dir',
            'ref_EMBL_files',
            'ref_bases',
            'ref_sequences',
            'ref_bases_assembled',
            'ref_sequences_assembled',
            'ref_sequences_assembled_ok',
            'ref_bases_assembler_missed',
            'assembly_bases',
            'assembly_bases_in_ref',
            'assembly_contigs',
            'assembly_contigs_hit_ref',
            'assembly_bases_reads_disagree',
            'assembly_sum_longest_match_each_segment',
            'cds_number',
            'cds_assembled',
            'cds_assembled_ok',
        ]
        self.stats = {shrubbery: -1 for shrubbery in self.stats_keys}
        self.stats_file_txt = output_prefix + '.stats.txt'
        self.stats_file_tsv = output_prefix + '.stats.tsv'

    def _set_assembly_fasta_data(self, fasta_filename):
        self.assembly_is_empty = os.path.getsize(fasta_filename) == 0
        if self.assembly_is_empty:
            return

        pyfastaq.tasks.to_fasta(fasta_filename, self.assembly_fasta, strip_after_first_whitespace=True)

        self.assembly_fasta_fai = self.assembly_fasta + '.fai'
        if not os.path.exists(self.assembly_fasta_fai):
            iva.common.syscall('samtools faidx ' + self.assembly_fasta)

        self.assembly_lengths = {}
        pyfastaq.tasks.lengths_from_fai(self.assembly_fasta_fai, self.assembly_lengths)

    def _set_ref_seq_data(self):
        assert self.embl_dir is not None
        self.ref_fasta = self.outprefix + '.reference.fa'
        self.ref_fasta_fai = self.ref_fasta + '.fai'
        self.ref_gff = self.outprefix + '.reference.gff'
        tmpdir = tempfile.mkdtemp(prefix='tmp.set_ref_seq_data.', dir=os.getcwd())
        extractor = iva.egg_extract.Extractor(os.path.abspath(os.path.join(os.path.dirname(iva.__file__), os.pardir)))
        embl2gff_egg = os.path.join('iva', 'ratt', 'embl2gff.pl')
        embl2gff = os.path.join(tmpdir, 'embl2gff.pl')
        extractor.copy_file(embl2gff_egg, embl2gff)
        os.chmod(embl2gff, stat.S_IRWXU)

        for embl_file in os.listdir(self.embl_dir):
            fa = os.path.join(tmpdir, embl_file + '.fa')
            gff = os.path.join(tmpdir, embl_file + '.gff')
            embl_full = os.path.join(self.embl_dir, embl_file)
            pyfastaq.tasks.to_fasta(embl_full, fa)
            iva.common.syscall(' '.join([embl2gff, embl_full, '>', gff]))

        iva.common.syscall('cat ' + tmpdir + '/*.gff > ' + self.ref_gff)
        iva.common.syscall('cat ' + tmpdir + '/*.fa > ' + self.ref_fasta)
        shutil.rmtree(tmpdir)
        self._set_ref_fa_data()

    def _set_ref_fa_data(self):
        self.ref_fasta_fai = self.ref_fasta + '.fai'
        iva.common.syscall('samtools faidx ' + self.ref_fasta)
        self.ref_ids = self._ids_in_order_from_fai(self.ref_fasta_fai)
        self.ref_lengths = {}
        pyfastaq.tasks.lengths_from_fai(self.ref_fasta_fai, self.ref_lengths)
        number_of_sequences = pyfastaq.tasks.count_sequences(self.ref_fasta)
        if number_of_sequences != len(self.ref_lengths):
            raise Error('At least one reference name has been used more than once.\nNames must be unique. Cannot continue')

        self.ref_length_offsets = {}
        offset = 0
        for seq in self.ref_ids:
            self.ref_length_offsets[seq] = offset
            offset += self.ref_lengths[seq]

    def _ids_in_order_from_fai(self, filename):
        ids = []
        f = pyfastaq.utils.open_file_read(filename)
        for line in f:
            ids.append(line.rstrip().split('\t')[0])
        pyfastaq.utils.close(f)
        return ids

    def _get_ref_cds_from_gff(self):
        f = pyfastaq.utils.open_file_read(self.ref_gff)
        coords = {}
        for line in f:
            # no annotation allowed after any fasta sequence. See
            # http://www.sequenceontology.org/gff3.shtml
            if line.rstrip() == '##FASTA':
                break
            elif line.startswith('#'):
                continue

            data = line.rstrip().split('\t')
            if data[2] == 'CDS':
                seqname = data[0]
                start = int(data[3]) - 1
                end = int(data[4]) - 1
                strand = data[6]
                if seqname not in coords:
                    coords[seqname] = []
                coords[seqname].append((pyfastaq.intervals.Interval(start, end), strand))

        pyfastaq.utils.close(f)
        for seqname in coords:
            coords[seqname].sort()

        return coords

    def _write_cds_seqs(self, cds_list, fa, f_out):
        for coords, strand in cds_list:
            seqname = fa.id + ':' + str(coords.start + 1) + '-' + str(coords.end + 1) + ':' + strand
            seq = pyfastaq.sequences.Fasta(seqname, fa.seq[coords.start:coords.end+1])
            if strand == '-':
                seq.revcomp()
            print(seq, file=f_out)
            assert seqname not in self.cds_assembly_stats
            self.cds_assembly_stats[seqname] = {
                'ref_name': fa.id,
                'ref_coords': coords,
                'strand': strand,
                'length_in_ref': len(coords),
                'assembled': False,
                'assembled_ok': False,
            }

    def _gff_and_fasta_to_cds(self):
        cds_coords = self._get_ref_cds_from_gff()
        f = pyfastaq.utils.open_file_write(self.ref_cds_fasta)
        seq_reader = pyfastaq.sequences.file_reader(self.ref_fasta)
        for seq in seq_reader:
            if seq.id in cds_coords:
                self._write_cds_seqs(cds_coords[seq.id], seq, f)
        pyfastaq.utils.close(f)

    def _map_cds_to_assembly(self):
        if not os.path.exists(self.ref_cds_fasta):
            self._gff_and_fasta_to_cds()
        if not self.assembly_is_empty:
            iva.mummer.run_nucmer(self.ref_cds_fasta, self.assembly_fasta, self.cds_nucmer_coords_in_assembly, min_length=self.nucmer_min_cds_hit_length, min_id=self.nucmer_min_cds_hit_id)

    def _mummer_coords_file_to_dict(self, filename):
        hits = {}
        for hit in iva.mummer.file_reader(filename):
            if hit.qry_name not in hits:
                hits[hit.qry_name] = []
            hits[hit.qry_name].append(copy.copy(hit))
        return hits

    def _has_orf(self, fa, start, end, min_length):
        subseq = pyfastaq.sequences.Fasta('seq', fa[start:end+1])
        orfs = subseq.all_orfs(min_length=min_length)
        return len(orfs) > 0

    def _calculate_cds_assembly_stats(self):
        if self.assembly_is_empty:
            return
        self._map_cds_to_assembly()
        hits = self._mummer_coords_file_to_dict(self.cds_nucmer_coords_in_assembly)
        contigs = {}
        pyfastaq.tasks.file_to_dict(self.assembly_fasta, contigs)
        for cds_name, hit_list in hits.items():
            self.cds_assembly_stats[cds_name]['number_of_contig_hits'] = len(hit_list)
            hit_coords = [x.qry_coords() for x in hit_list]
            pyfastaq.intervals.merge_overlapping_in_list(hit_coords)
            bases_assembled = pyfastaq.intervals.length_sum_from_list(hit_coords)
            self.cds_assembly_stats[cds_name]['bases_assembled'] = bases_assembled
            self.cds_assembly_stats[cds_name]['assembled'] = 0.9 <= bases_assembled / self.cds_assembly_stats[cds_name]['length_in_ref'] <= 1.1

            if len(hit_list) == 1:
                hit = hit_list[0]
                contig_coords = hit.ref_coords()
                has_orf = self._has_orf(contigs[hit.ref_name], contig_coords.start, contig_coords.end, 0.9 * self.cds_assembly_stats[cds_name]['length_in_ref'])
                self.cds_assembly_stats[cds_name]['assembled_ok'] = has_orf
            else:
                self.cds_assembly_stats[cds_name]['assembled_ok'] = False

    def _get_contig_hits_to_reference(self):
        iva.mummer.run_nucmer(self.assembly_fasta, self.ref_fasta, self.assembly_vs_ref_coords, min_id=self.nucmer_min_ctg_hit_id, min_length=self.nucmer_min_ctg_hit_length, breaklen=500)
        self.assembly_vs_ref_mummer_hits = self._mummer_coords_file_to_dict(self.assembly_vs_ref_coords)

    def _write_fasta_contigs_hit_ref(self):
        contigs = {}
        pyfastaq.tasks.file_to_dict(self.assembly_fasta, contigs)
        f = pyfastaq.utils.open_file_write(self.fasta_assembly_contigs_hit_ref)
        for qry_name in sorted(self.assembly_vs_ref_mummer_hits):
            print(contigs[qry_name], file=f)
        pyfastaq.utils.close(f)

    def _write_fasta_contigs_not_hit_ref(self):
        seq_reader = pyfastaq.sequences.file_reader(self.assembly_fasta)
        f = pyfastaq.utils.open_file_write(self.fasta_assembly_contigs_not_hit_ref)
        for seq in seq_reader:
            if seq.id not in self.assembly_vs_ref_mummer_hits:
                print(seq, file=f)
        pyfastaq.utils.close(f)

    def _hash_nucmer_hits_by_ref(self, hits):
        d = {}
        for l in hits.values():
            for hit in l:
                if hit.ref_name not in d:
                    d[hit.ref_name] = []
                d[hit.ref_name].append(copy.copy(hit))
        return d

    def _calculate_refseq_assembly_stats(self):
        if self.assembly_is_empty:
            refhits = {}
        else:
            refhits = self._hash_nucmer_hits_by_ref(self.assembly_vs_ref_mummer_hits)

        for name in self.ref_ids:
            assert name not in self.refseq_assembly_stats
            if name in refhits:
                hits = refhits[name]
                coords = [hit.ref_coords() for hit in hits]
                pyfastaq.intervals.merge_overlapping_in_list(coords)
                self.refseq_assembly_stats[name] = {
                    'hits': len(hits),
                    'bases_assembled': pyfastaq.intervals.length_sum_from_list(coords),
                    'assembled': 0.9 <= pyfastaq.intervals.length_sum_from_list(coords) / self.ref_lengths[name],
                    'assembled_ok': len(hits) == 1 and 0.9 <= hits[0].hit_length_ref / self.ref_lengths[name] <= 1.1 and 0.9 <= hits[0].qry_length / hits[0].hit_length_qry,
                    'longest_matching_contig': self._longest_matching_contig(refhits, name),
                }
            else:
                self.refseq_assembly_stats[name] = {
                    'hits': 0,
                    'bases_assembled': 0,
                    'assembled': False,
                    'assembled_ok': False,
                    'longest_matching_contig': 0,
                }

    def _invert_list(self, coords, seq_length):
        if len(coords) == 0:
            return[pyfastaq.intervals.Interval(0, seq_length - 1)]

        not_covered = []

        if coords[0].start != 0:
            not_covered.append(pyfastaq.intervals.Interval(0, coords[0].start - 1))

        for i in range(len(coords) - 1):
            not_covered.append(pyfastaq.intervals.Interval(coords[i].end + 1, coords[i+1].start - 1))

        if coords[-1].end < seq_length - 1:
            not_covered.append(pyfastaq.intervals.Interval(coords[-1].end + 1, seq_length - 1))

        return not_covered

    def _calculate_ref_positions_covered_by_contigs(self):
        if self.assembly_is_empty:
            self.ref_pos_not_covered_by_contigs = {}
            for seq, lngth in self.ref_lengths.items():
                self.ref_pos_not_covered_by_contigs[seq] = [pyfastaq.intervals.Interval(0, lngth - 1)]
            return

        for seq in self.assembly_vs_ref_mummer_hits:
            for hit in self.assembly_vs_ref_mummer_hits[seq]:
                if hit.ref_name not in self.ref_pos_covered_by_contigs:
                    self.ref_pos_covered_by_contigs[hit.ref_name] = []
                self.ref_pos_covered_by_contigs[hit.ref_name].append(hit.ref_coords())

        for coords_list in self.ref_pos_covered_by_contigs.values():
            pyfastaq.intervals.merge_overlapping_in_list(coords_list)

        for seq in self.ref_ids:
            if seq in self.ref_pos_covered_by_contigs:
                l = self.ref_pos_covered_by_contigs[seq]
            else:
                l = []
            self.ref_pos_not_covered_by_contigs[seq] = self._invert_list(l, self.ref_lengths[seq])

    def _get_overlapping_qry_hits(self, hits, hit):
        overlapping = []
        hit_coords = hit.qry_coords()

        for test_hit in hits:
            if test_hit != hit:
                test_coords = test_hit.qry_coords()
                if test_coords.intersects(hit_coords):
                    overlapping.append(test_hit)

        return overlapping

    def _get_unique_and_repetitive_from_contig_hits(self, hits):
        unique = []
        repetitive = []
        if len(hits) == 0:
            return unique, repetitive

        for hit in hits:
            if len(self._get_overlapping_qry_hits(hits, hit)):
                repetitive.append(hit)
            else:
                unique.append(hit)

        return unique, repetitive

    def _longest_matching_contig(self, hits, refname):
        if len(hits) == 0:
            return 0
        hits_to_refname = [x for x in hits[refname] if x.ref_name == refname and x.qry_length <= 2 * x.ref_length]
        if len(hits_to_refname) == 0:
            return 0
        else:
            return max([x.qry_length for x in hits_to_refname])

    def _calculate_incorrect_assembly_bases(self):
        if self.assembly_is_empty:
            self.incorrect_assembly_bases = {}
        else:
            self.incorrect_assembly_bases = iva.mapping.find_incorrect_ref_bases(self.assembly_bam, self.assembly_fasta)

    def _contig_placement_in_reference(self, hits):
        unique_hits, repetitive_hits = self._get_unique_and_repetitive_from_contig_hits(hits)
        placement = [(x.qry_coords(), x.ref_name, x.ref_coords(), x.on_same_strand(), False) for x in unique_hits]
        placement += [(x.qry_coords(), x.ref_name, x.ref_coords(), x.on_same_strand(), True) for x in repetitive_hits]
        placement.sort()
        return placement

    def _calculate_contig_placement(self):
        if self.assembly_is_empty:
            return
        self._get_contig_hits_to_reference()
        self.contig_placement = {qry_name: self._contig_placement_in_reference(self.assembly_vs_ref_mummer_hits[qry_name]) for qry_name in self.assembly_vs_ref_mummer_hits}

    def _get_R_plot_contig_order_from_contig_placement(self):
        contig_positions = []
        for qryname, coords_list in  self.contig_placement.items():
            for qry_coords, refname, ref_coords, same_strand, repetitive in coords_list:
                offset = self.ref_length_offsets[refname]
                ref_coords = pyfastaq.intervals.Interval(ref_coords.start + offset, ref_coords.end + offset)
                contig_positions.append((ref_coords, qry_coords, same_strand, repetitive, qryname))

        contig_positions.sort()
        names = set()
        contig_names = []
        for name in [x[4] for x in contig_positions]:
            if name not in names:
                contig_names.append(name)
                names.add(name)
        return contig_names

    def _map_reads_to_assembly(self):
        if not self.assembly_is_empty:
            iva.mapping.map_reads(self.reads_fwd, self.reads_rev, self.assembly_fasta, self.assembly_bam[:-4], sort=True, threads=self.threads, index_k=self.smalt_k, index_s=self.smalt_s, minid=self.smalt_id, extra_smalt_map_ops='-x')
            os.unlink(self.assembly_bam[:-4] + '.unsorted.bam')

    def _write_ref_info(self, filename):
        assert self.embl_dir is not None
        files = sorted(os.listdir(self.embl_dir))
        f = pyfastaq.utils.open_file_write(filename)
        print('EMBL_directory', self.embl_dir, sep='\t', file=f)
        print('Files', '\t'.join(files), sep='\t', file=f)
        pyfastaq.utils.close(f)

    def _choose_reference_genome(self):
        if self.embl_dir is None:
            assert self.ref_db is not None
            assert os.path.exists(self.assembly_bam)
            tmp_reads = self.outprefix + '.tmp.subsample.reads.fastq'
            iva.mapping.subsample_bam(self.assembly_bam, tmp_reads, coverage=40)
            db = iva.kraken.Database(self.ref_db, threads=self.threads, preload=self.kraken_preload)
            self.embl_dir = db.choose_reference(tmp_reads, self.kraken_prefix)
            os.unlink(tmp_reads)
            if self.embl_dir is None:
                raise Error('Unable to determine reference genome automatically. Cannot continue')
        else:
            self.embl_dir = os.path.abspath(self.embl_dir)

        self.embl_files = sorted(os.listdir(self.embl_dir))

    def _make_act_files(self):
        if self.assembly_is_empty or not self.blast_for_act:
            return

        iva.qc_external.run_blastn_and_write_act_script(self.assembly_fasta, self.ref_fasta, self.blast_out, self.act_script)

    def _map_reads_to_reference(self):
        assert os.path.exists(self.ref_fasta)
        iva.mapping.map_reads(self.reads_fwd, self.reads_rev, self.ref_fasta, self.ref_bam[:-4], sort=True, threads=self.threads, index_k=self.smalt_k, index_s=self.smalt_s, minid=self.smalt_id, extra_smalt_map_ops='-x')
        os.unlink(self.ref_bam[:-4] + '.unsorted.bam')

    def _calculate_ref_read_coverage(self):
        if not os.path.exists(self.ref_bam):
            self._map_reads_to_reference()
        for seq in self.ref_ids:
            assert seq not in self.ref_coverage_fwd
            self.ref_coverage_fwd[seq] = iva.mapping.get_bam_region_coverage(self.ref_bam, seq, self.ref_lengths[seq])
            assert seq not in self.ref_coverage_rev
            self.ref_coverage_rev[seq] = iva.mapping.get_bam_region_coverage(self.ref_bam, seq, self.ref_lengths[seq], rev=True)

    def _coverage_list_to_low_cov_intervals(self, l):
        bad_intervals = []
        start = None
        cov_bad = False

        for i in range(len(l)):
            cov_bad = l[i] < self.min_ref_cov
            if cov_bad:
                if start is None:
                    start = i
            else:
                if start is not None:
                    bad_intervals.append(pyfastaq.intervals.Interval(start, i-1))
                start = None

        if cov_bad and start is not None:
            bad_intervals.append(pyfastaq.intervals.Interval(start, i))
        return bad_intervals

    def _calculate_ref_read_region_coverage(self):
        assert len(self.ref_coverage_fwd)
        assert len(self.ref_coverage_rev)
        for seq in self.ref_ids:
            self.low_cov_ref_regions_fwd[seq] = self._coverage_list_to_low_cov_intervals(self.ref_coverage_fwd[seq])
            self.low_cov_ref_regions_rev[seq] = self._coverage_list_to_low_cov_intervals(self.ref_coverage_rev[seq])
            fwd_ok = self._invert_list(self.low_cov_ref_regions_fwd[seq], self.ref_lengths[seq])
            rev_ok = self._invert_list(self.low_cov_ref_regions_rev[seq], self.ref_lengths[seq])
            self.ok_cov_ref_regions[seq] = pyfastaq.intervals.intersection(fwd_ok, rev_ok)
            self.low_cov_ref_regions[seq] = pyfastaq.intervals.intersection(self.low_cov_ref_regions_fwd[seq], self.low_cov_ref_regions_rev[seq])

    def _write_ref_coverage_to_files_for_R(self, outprefix):
        assert len(self.ref_coverage_fwd)
        assert len(self.ref_coverage_rev)
        def list_to_file(d, fname):
            f = pyfastaq.utils.open_file_write(fname)
            for refname in self.ref_ids:
                for x in d[refname]:
                    print(x, file=f)
            pyfastaq.utils.close(f)
        list_to_file(self.ref_coverage_fwd, outprefix + '.fwd')
        list_to_file(self.ref_coverage_rev, outprefix + '.rev')

    def _cov_to_R_string(self, intervals, colour, x_offset, y_position, contig_height):
        s = ''
        for interval in intervals:
            s += 'rect(' + \
                 str(interval.start + x_offset) + ', ' + \
                 str(y_position - 0.5 * contig_height) + ', ' + \
                 str(interval.end + x_offset) + ', ' + \
                 str(y_position + 0.5 * contig_height) + ', ' + \
                 'col="' + colour + '", ' + \
                 'border=NA)\n'
        return s

    def _calculate_should_have_assembled(self):
        for name in self.ref_ids:
            if name in self.ref_pos_covered_by_contigs:
                l = self.ref_pos_covered_by_contigs[name]
            else:
                l = []
            self.should_have_assembled[name] = pyfastaq.intervals.intersection(self._invert_list(l, self.ref_lengths[name]), self.ok_cov_ref_regions[name])

    def _calculate_gage_stats(self):
        if self.assembly_is_empty:
            self.gage_stats = iva.qc_external.dummy_gage_stats()
            self.gage_stats['Missing Reference Bases'] = sum(self.ref_lengths.values())
        else:
            self.gage_stats = iva.qc_external.run_gage(self.ref_fasta, self.assembly_fasta, self.gage_outdir, nucmer_minid=self.gage_nucmer_minid, clean=self.clean)

    def _calculate_ratt_stats(self):
        if self.assembly_is_empty:
            self.ratt_stats = iva.qc_external.dummy_ratt_stats()
        else:
            self.ratt_stats = iva.qc_external.run_ratt(self.embl_dir, self.assembly_fasta, self.ratt_outdir, config_file=self.ratt_config, clean=self.clean)

    def _do_calculations(self):
        self._map_reads_to_assembly()
        self._choose_reference_genome()
        self._set_ref_seq_data()
        self._make_act_files()
        self._map_reads_to_reference()
        self._calculate_incorrect_assembly_bases()
        self._calculate_contig_placement()
        self._write_fasta_contigs_hit_ref()
        self._write_fasta_contigs_not_hit_ref()
        self._calculate_ref_read_coverage()
        self._calculate_ref_read_region_coverage()
        self._calculate_ref_positions_covered_by_contigs()
        self._calculate_should_have_assembled()
        self._calculate_cds_assembly_stats()
        self._calculate_refseq_assembly_stats()
        self._calculate_gage_stats()
        self._calculate_ratt_stats()
        self._calculate_stats()

    def _contigs_and_bases_that_hit_ref(self):
        total_bases = 0
        for name in self.assembly_vs_ref_mummer_hits:
            coords = [x.qry_coords() for x in self.assembly_vs_ref_mummer_hits[name]]
            pyfastaq.intervals.merge_overlapping_in_list(coords)
            total_bases += pyfastaq.intervals.length_sum_from_list(coords)
        return total_bases, len(self.assembly_vs_ref_mummer_hits)

    def _calculate_stats(self):
        self.stats['ref_EMBL_dir'] = self.embl_dir
        self.stats['ref_EMBL_files'] = ' '.join(self.embl_files)
        self.stats['ref_bases'] = sum(self.ref_lengths.values())
        self.stats['ref_sequences'] = len(self.ref_lengths)
        self.stats['ref_bases_assembled'] = sum([pyfastaq.intervals.length_sum_from_list(l) for l in list(self.ref_pos_covered_by_contigs.values())])
        self.stats['ref_sequences_assembled'] = len([1 for x in self.refseq_assembly_stats.values() if x['assembled']])
        self.stats['ref_sequences_assembled_ok'] = len([1 for x in self.refseq_assembly_stats.values() if x['assembled_ok']])
        self.stats['ref_bases_assembler_missed'] = sum([pyfastaq.intervals.length_sum_from_list(l) for l in list(self.should_have_assembled.values())])
        self.stats['assembly_bases'] = sum(self.assembly_lengths.values())
        self.stats['assembly_contigs'] = len(self.assembly_lengths)
        self.stats['assembly_bases_in_ref'], self.stats['assembly_contigs_hit_ref'] = self._contigs_and_bases_that_hit_ref()
        self.stats['assembly_bases_reads_disagree'] = sum([len(x) for x in self.incorrect_assembly_bases.values()])
        self.stats['assembly_sum_longest_match_each_segment'] = sum([x['longest_matching_contig'] for x in self.refseq_assembly_stats.values()])
        self.stats['cds_number'] = len(self.cds_assembly_stats)
        self.stats['cds_assembled'] = len([1 for x in self.cds_assembly_stats.values() if x['assembled']])
        self.stats['cds_assembled_ok'] = len([1 for x in self.cds_assembly_stats.values() if x['assembled_ok']])

    def _write_stats_txt(self):
        f = pyfastaq.utils.open_file_write(self.stats_file_txt)
        for stat in self.stats_keys:
            print(stat, self.stats[stat], sep='\t', file=f)
        for stat in iva.qc_external.gage_stats:
            print('gage_' + stat.replace(' ', '_'), self.gage_stats[stat], sep='\t', file=f)
        for stat in iva.qc_external.ratt_stats:
            print('ratt_' + stat.replace(' ', '_'), self.ratt_stats[stat], sep='\t', file=f)

        pyfastaq.utils.close(f)

    def _write_stats_tsv(self):
        f = pyfastaq.utils.open_file_write(self.stats_file_tsv)
        print('\t'.join([x.replace(' ', '_') for x in self.stats_keys]), end='\t', file=f)
        print('\t'.join(['gage_' + x.replace(' ', '_') for x in iva.qc_external.gage_stats]), end='\t', file=f)
        print('\t'.join(['ratt_' + x.replace(' ', '_') for x in iva.qc_external.ratt_stats]), file=f)
        print('\t'.join([str(self.stats[x]) for x in self.stats_keys]),
              '\t'.join([str(self.gage_stats[x]) for x in iva.qc_external.gage_stats]),
              '\t'.join([str(self.ratt_stats[x]) for x in iva.qc_external.ratt_stats]),
              sep='\t', file=f)
        pyfastaq.utils.close(f)

    def _write_stats_files(self):
        self._write_stats_txt()
        self._write_stats_tsv()

    def _make_R_plots(self):
        outprefix = self.outprefix + '.contig_placement'
        contig_names = self._get_R_plot_contig_order_from_contig_placement()
        number_of_contigs = len(contig_names)
        ref_length = sum(self.ref_lengths.values())
        r_script = outprefix + '.R'
        f = pyfastaq.utils.open_file_write(r_script)
        contig_height = 0.8
        vertical_lines = ''
        if len(self.ref_ids) > 0:
            for name in self.ref_ids:
                x_position = self.ref_length_offsets[name]
                if x_position > 0:
                    vertical_lines += '\n' + 'abline(v=' + str(x_position) + ', col="gray")'

        print('pdf(file="', outprefix, '.pdf")', sep='', file=f)
        print('layout(matrix(c(1,2,3), 3, 1, byrow = TRUE), heights=c(3,1.2,1.2))', file=f)
        print('par(mar=c(1, 4, 2, 0.5))', file=f)

        # ---------- contig layout plot ------------------------
        print('plot(-100, type="n", xlim=c(0,', ref_length, '), ylim=c(0, ', number_of_contigs, '), yaxt="n", ylab="", xlab="")', sep='', file=f)
        print('title("', self.contig_layout_plot_title, '", ylab="Contigs")', sep='', file=f)
        print(vertical_lines, file=f)

        if number_of_contigs > 0:
            print('contig_names=c("', '", "'.join(contig_names), '")', sep='', file=f)
            print('axis(2, at=c(1:', number_of_contigs, '), labels=contig_names, las=2, cex.axis=0.3)', sep='', file=f)

            for i in range(len(contig_names)):
                contig_name = contig_names[i]
                y_centre = i + 1
                contig_positions = self.contig_placement[contig_name]
                for contig_coords, ref_name, ref_coords, same_strand, repetitive in contig_positions:
                    offset = self.ref_length_offsets[ref_name]
                    if repetitive:
                        colour = "red"
                    else:
                        colour = "blue"

                    if same_strand:
                         colour = "dark" + colour

                    print('rect(', ref_coords.start + offset, ',',
                            y_centre - 0.5 * contig_height, ',',
                            ref_coords.end + offset, ',',
                            y_centre + 0.5 * contig_height, ',',
                            'col="', colour, '")', sep='', file=f)

        # ----------- read coverage heatmap ---------------------
        print('par(mar=c(0, 4, 1, 0.5))', file=f)
        print('plot(-100, type="n", xlim=c(0,', ref_length, '), ylim=c(0, ', 3, '), xaxt="n", yaxt="n", ylab="Contig/Read coverage OK", xlab="", frame.plot=F)', sep='', file=f)
        print('axis(2, at=c(1,2), labels=c("Reads", "Contigs"), las=2, cex.axis=0.6)', sep='', file=f)

        for name in self.ref_ids:
            offset = self.ref_length_offsets[name]
            print(self._cov_to_R_string(self.ok_cov_ref_regions[name], 'black', offset, 1.3, 0.25), file=f)
            print(self._cov_to_R_string(self.low_cov_ref_regions_fwd[name], 'red', offset, 1, 0.25), file=f)
            print(self._cov_to_R_string(self.low_cov_ref_regions_rev[name], 'red', offset, 0.7, 0.25), file=f)

            if name in self.ref_pos_covered_by_contigs:
                print(self._cov_to_R_string(self.ref_pos_covered_by_contigs[name], 'black', offset, 2.3, 0.25), file=f)

            if name in self.should_have_assembled:
                print(self._cov_to_R_string(self.should_have_assembled[name], 'red', offset, 1.7, 0.25), file=f)

            print(self._cov_to_R_string(self.ref_pos_not_covered_by_contigs[name], 'black', offset, 2, 0.25), file=f)
        print(vertical_lines, file=f)

        # ----------- read depth on reference plot --------------
        self._write_ref_coverage_to_files_for_R(self.outprefix + '.read_coverage_on_ref')
        print('fwd_ref_cov = scan("', self.outprefix + '.read_coverage_on_ref.fwd', '")', sep='', file=f)
        print('rev_ref_cov = scan("', self.outprefix + '.read_coverage_on_ref.rev', '")', sep='', file=f)
        print('par(mar=c(5, 4, 0, 0.5))', file=f)
        print('plot(fwd_ref_cov, type="l", xlim=c(0, length(fwd_ref_cov) + 1), ylim=c(-max(rev_ref_cov), max(fwd_ref_cov)), col="blue", frame.plot=F, ylab="Read depth", xlab="Position in reference")', file=f)
        print('lines(-rev_ref_cov, col="blue")', file=f)
        print('abline(h=0, lty=2)', file=f)
        print(vertical_lines, file=f)

        print('dev.off()', file=f)
        pyfastaq.utils.close(f)
        iva.common.syscall('R CMD BATCH ' + r_script)

    def _clean(self):
        for fname in self.files_to_clean:
            os.unlink(fname)

    def run(self):
        self._do_calculations()
        self._make_R_plots()
        self._write_stats_files()
        self._clean()


#for content from cloud

from regex import D
import requests
import csv

from xml.dom.minidom import parseString
url = 'http://api.irishrail.ie/realtime/realtime.asmx/getCurrentTrainsXML'

page = requests.get(url)
doc = parseString(page.content)

retrieveTags=['TrainStatus',
'TrainLatitude',
'TrainLongitude',
'TrainCode',
'TrainDate',
'PublicMessage',
'Direction'
]

objTrainPositionsNodes = doc.getElementsByTagName("objTrainPositions")
for objTrainPositionsNode in objTrainPositionsNodes:
    traincodenode = objTrainPositionsNode.getElementsByTagName("TrainCode").item(0)

    traincode = traincodenode.firstChild.nodeValue.strip()

    print (traincode)


from typing import Any, Sequence

from sqlalchemy import select, or_, and_

from entity import Coordinate, City, Road, Place
from .repository import Repository

class CoordinateRepository(Repository):
    def __init__(self):
        super().__init__(Coordinate)

    def find_all(self) -> Sequence[Coordinate]:
        return super().find_all()

    def find_by(self, condition: Any) -> Sequence[Coordinate]:
        return super().find_by(condition)

    def find_by_id(self, id_: int) -> Coordinate | None:
        return super().find_by_id(id_)

    def find_by_city_left_joined_place_where_place_coordinate_is_null(self, city: City) -> Sequence[Coordinate]:
        statement = (
            select(self._entity_type).
            join(Place, Place.coordinate_id == Coordinate.id, isouter=True).
            where(and_(Coordinate.city == city, Place.coordinate == None))
        )
        return super()._fetch_all(statement)

    def find_by_city_joined_road_ordered_by_id(self, city: City) -> Sequence[tuple[Coordinate, Road]]:
        statement = (
            select(self._entity_type, Road).
            join(Road, Road.point_0_id == Coordinate.id).
            where(Coordinate.city == city).
            order_by(Coordinate.id)
        )
        return super()._fetch_all(statement, scalar=False)

    def find_by_city_joined_road_with_repetitions_ordered_by_id(self, city: City) -> Sequence[tuple[Coordinate, Road]]:
        statement = (
            select(self._entity_type, Road).
            join(Road, or_(Road.point_0_id == Coordinate.id, Road.point_1_id == Coordinate.id)).
            where(Coordinate.city == city).
            order_by(Coordinate.id)
        )
        return super()._fetch_all(statement, scalar=False)


import h5py
import numpy as np

from utils.registrator import Registrator

def progress(i, total):

    fraction = i / total
    bar_length = 50
    bar_fraction = int(fraction*bar_length)
    bar = '#'*bar_fraction + '='*(bar_length-bar_fraction)
    print(f'\r|{bar}| {round(fraction*100, 2)}%\t', end='')

def preprocess(sample, factor=0.5):
    sample = sample/sample.max()
    mean = sample[sample>0].mean()
    std = sample[sample>0].std()
    level = mean - factor*std
    sample = sample.clip(level, 1.0) - level
    sample = sample/sample.max()
    sample[0:8] = 0
    return sample

storage_dir = '/media/network/DL_PC/Datasets/oct_quality_validation/'
lq = h5py.File(storage_dir+'low.hdf5', 'r')
hq = h5py.File(storage_dir+'high.hdf5', 'r')
output = h5py.File(storage_dir+'registered.hdf5', 'a')

registrator = Registrator(order=3, offset=0)

lq_keys, hq_keys = list(lq.keys()), list(hq.keys())
n_samples = len(lq_keys)

try:
    for i, (lq_key, hq_key) in enumerate(zip(lq_keys, hq_keys)):

        subject, template = preprocess(lq[lq_key].value), preprocess(hq[hq_key].value)
        try:
            registered, transformation, difference = registrator.register_single(template, subject)
            transformation.pop('timg')
        except (IndexError, ValueError):
            print(f'\nsomething went wrong in the registration of {lq_key}')
            registered = np.zeros_like(subject)
            difference = 1
            transformation = {}
        output[lq_key] = registered
        for key, value in transformation.items():
            output[lq_key].attrs[key] = value
        output[lq_key].attrs['difference'] = difference
        # output.flush()
        progress(i, n_samples)
finally:
    output.close()
    lq.close()
    hq.close()

"""
something went wrong in the registration of 17182-624913-00
something went wrong in the registration of 17182-624913-01
something went wrong in the registration of 17182-624913-02
something went wrong in the registration of 17802-626829-14
something went wrong in the registration of 17802-626829-15
something went wrong in the registration of 7681-624748-00
something went wrong in the registration of 7681-624748-01
something went wrong in the registration of 7681-624748-10
something went wrong in the registration of 7681-624748-11
something went wrong in the registration of 7681-624748-13
"""


import time
import traceback
from datetime import timedelta, datetime
from multiprocessing import Process
from library.logger import getLogger
import requests
from constants.config import config
from azure.storage.blob import BlobServiceClient
from azure.storage.queue import QueueServiceClient
from azure.core import exceptions as AzureExceptions
import library.db as db
import json
import library.utils as utils

logger = getLogger("validate")

def process_hash_list(document_datasets):

    conn = db.getDirectConnection()
    now = datetime.now()

    for file_data in document_datasets:
        try:
            file_hash = file_data[0]
            downloaded = file_data[1]
            file_id = file_data[2]
            file_url = file_data[3]
            publisher = file_data[4]
            publisher_name = file_data[5]
            file_schema_valid = file_data[6]
            publisher_black_flag = file_data[7] is not None

            if file_schema_valid == False and downloaded > (now - timedelta(hours=config['VALIDATION']['SAFETY_CHECK_PERIOD'])):
                logger.info(
                    f"Skipping Schema Invalid file for Full Validation until {config['VALIDATION']['SAFETY_CHECK_PERIOD']}hrs after download: {downloaded.isoformat()} for hash: {file_hash} and id: {file_id}")
                continue

            if file_schema_valid == False and publisher_black_flag == True:
                logger.info(
                    f"Skipping Schema Invalid file for Full Validation since publisher: {publisher} is black flagged for hash: {file_hash} and id: {file_id}")
                continue

            blob_name = file_hash + '.xml'

            blob_service_client = BlobServiceClient.from_connection_string(
                config['STORAGE_CONNECTION_STR'])
            blob_client = blob_service_client.get_blob_client(
                container=config['SOURCE_CONTAINER_NAME'], blob=blob_name)

            downloader = blob_client.download_blob()

            try:
                payload = utils.get_text_from_blob(downloader, file_hash)
            except:
                logger.warning(
                    f"Could not identify charset for hash: {file_hash} and id: {file_id}")
                continue

            if file_schema_valid is None:
                logger.info(
                    f"Schema Validating file hash: {file_hash} and id: {file_id}")
                schema_headers = {config['VALIDATION']['SCHEMA_VALIDATION_KEY_NAME']: config['VALIDATION']['SCHEMA_VALIDATION_KEY_VALUE']}
                schema_response = requests.post(
                    config['VALIDATION']['SCHEMA_VALIDATION_URL'], data=payload.encode('utf-8'), headers=schema_headers)
                db.updateValidationRequestDate(conn, file_id)

                if schema_response.status_code != 200:
                    if schema_response.status_code >= 400 and schema_response.status_code < 500:  # client errors
                        # log in db and 'continue' to break out of for loop for this file
                        db.updateValidationError(
                            conn, file_id, schema_response.status_code)
                        logger.warning(
                            f"Schema Validator reports Client Error HTTP {schema_response.status_code} for hash: {file_hash} and id: {file_id}")
                        continue
                    elif schema_response.status_code >= 500:  # server errors
                        # log in db and 'continue' to break out of for loop for this file
                        db.updateValidationError(
                            conn, file_id, schema_response.status_code)
                        logger.warning(
                            f"Schema Validator reports Server Error HTTP {schema_response.status_code} for hash: {file_hash} and id: {file_id}")
                        continue
                    else:
                        logger.error(
                            f"Schema Validator reports HTTP {schema_response.status_code} for hash: {file_hash} and id: {file_id}")
                try:
                    body = schema_response.json()
                    if body['valid'] == True or body['valid'] == False:
                        db.updateDocumentSchemaValidationStatus(
                            conn, file_id, body['valid'])
                        file_schema_valid = body['valid']
                    else:
                        raise
                except:
                    logger.error(
                        f"Unexpected response body from Schema validator for hash: {file_hash} and id: {file_id}")
                    continue

            if file_schema_valid == False and downloaded > (now - timedelta(hours=config['VALIDATION']['SAFETY_CHECK_PERIOD'])):
                logger.info(
                    f"Skipping Schema Invalid file for Full Validation until {config['VALIDATION']['SAFETY_CHECK_PERIOD']}hrs after download: {downloaded.isoformat()} for hash: {file_hash} and id: {file_id}")
                continue

            if file_schema_valid == False and publisher_black_flag == True:
                logger.info(
                    f"Skipping Schema Invalid file for Full Validation since publisher: {publisher} is black flagged for hash: {file_hash} and id: {file_id}")
                continue

            logger.info(
                f"Full Validating file hash: {file_hash} and id: {file_id}")

            full_headers = {config['VALIDATION']['FULL_VALIDATION_KEY_NAME']: config['VALIDATION']['FULL_VALIDATION_KEY_VALUE']}

            full_url = config['VALIDATION']['FULL_VALIDATION_URL']

            # only need meta=true for invalid files to "clean" them later
            if file_schema_valid == False:
                full_url += '?meta=true'
            full_response = requests.post(
                full_url, data=payload.encode('utf-8'), headers=full_headers)
            db.updateValidationRequestDate(conn, file_id)

            if full_response.status_code != 200:
                # 'expected' error codes returned from Validator
                if full_response.status_code == 400 or full_response.status_code == 413 or full_response.status_code == 422:
                    # log db and move on to save the validation report
                    db.updateValidationError(
                        conn, file_id, full_response.status_code)
                elif full_response.status_code >= 400 and full_response.status_code < 500:  # unexpected client errors
                    # log in db and 'continue' to break out of for loop for this file
                    db.updateValidationError(
                        conn, file_id, full_response.status_code)
                    logger.warning(
                        f"Full Validator reports Client Error HTTP {full_response.status_code} for hash: {file_hash} and id: {file_id}")
                    continue
                elif full_response.status_code >= 500:  # server errors
                    # log in db and 'continue' to break out of for loop for this file
                    db.updateValidationError(
                        conn, file_id, full_response.status_code)
                    logger.warning(
                        f"Full Validator reports Server Error HTTP {full_response.status_code} for hash: {file_hash} and id: {file_id}")
                    continue
                else:
                    logger.error(
                        f"Full Validator reports HTTP {full_response.status_code} for hash: {file_hash} and id: {file_id}")

            report = full_response.json()

            state = report.get('valid', None)

            db.updateValidationState(
                conn, file_id, file_hash, file_url, publisher, state, json.dumps(report), publisher_name)

        except (AzureExceptions.ResourceNotFoundError) as e:
            logger.warning(
                f"Blob not found for hash: {file_hash} and id: {file_id} updating as Not Downloaded for the refresher to pick up.")
            db.updateFileAsNotDownloaded(conn, file_id)
        except Exception as e:
            logger.error(f"ERROR with validating {file_hash}")
            print(traceback.format_exc())
            if hasattr(e, 'message'):
                logger.error(e.message)
            if hasattr(e, 'msg'):
                logger.error(e.msg)
            try:
                logger.warning(e.args[0])
            except:
                pass

    conn.close()

def validate():
    logger.info("Starting validation...")

    conn = db.getDirectConnection()

    file_hashes = db.getUnvalidatedDatasets(conn)

    if config['VALIDATION']['PARALLEL_PROCESSES'] == 1:
        logger.info(
            f"Processing {len(file_hashes)} IATI files in a single process for validation")
        process_hash_list(file_hashes)
    else:
        chunked_hash_lists = list(utils.chunk_list(
            file_hashes, config['VALIDATION']['PARALLEL_PROCESSES']))

        processes = []

        logger.info(
            f"Processing {len(file_hashes)} IATI files in a maximum of {config['VALIDATION']['PARALLEL_PROCESSES']} parallel processes for validation")

        for chunk in chunked_hash_lists:
            if len(chunk) == 0:
                continue
            process = Process(target=process_hash_list, args=(chunk,))
            process.start()
            processes.append(process)

        for process in processes:
            process.join()

    conn.close()
    logger.info("Finished validation.")

def safety_check():
    conn = db.getDirectConnection()

    logger.info(f"Starting validation safety check / publisher black flag check")

    try:
        queue_service_client = QueueServiceClient.from_connection_string(
            config['STORAGE_CONNECTION_STR'])
        queue_client = queue_service_client.get_queue_client(
            "publisher-black-flag-remove")

        messages = queue_client.receive_messages()

        for message in messages:
            try:
                logger.info(
                    f"Received message to remove black flag for publisher id: {message.content}")
                db.removeBlackFlag(conn, message.content)
                logger.info(f"Dequeueing message: {message.content}")
                queue_client.delete_message(message.id, message.pop_receipt)
            except Exception as e:
                logger.warning(
                    f"Could not process message with id: {message.id} for publisher id: {message.content}")
                continue
    except Exception as e:
        logger.warning(f"Failed to process removal of publisher black flags")

    db.blackFlagDubiousPublishers(
        conn, config['VALIDATION']['SAFETY_CHECK_THRESHOLD'], config['VALIDATION']['SAFETY_CHECK_PERIOD'])

    black_flags = db.getUnnotifiedBlackFlags(conn)

    for black_flag in black_flags:
        org_id = black_flag[0]

        notification = {
            "type": "NEW_BLACK_FLAG",
            "data": {
                "publisherId": org_id,
                "reason": f"Over {config['VALIDATION']['SAFETY_CHECK_THRESHOLD']} critical documents in the last {config['VALIDATION']['SAFETY_CHECK_PERIOD']}) hours."
            }
        }
        headers = {
            config['NOTIFICATION_KEY_NAME']: config['NOTIFICATION_KEY_VALUE'],
            'Content-Type': 'application/json'
        }

        try:
            response = requests.post(
                config['NOTIFICATION_URL'], data=json.dumps(notification), headers=headers)
        except Exception as e:
            logger.warning(
                f"Could not notify Black Flag for publisher id: {org_id}")
            continue

        if response.status_code != 200:
            logger.warning(
                f"Could not notify Black Flag for publisher id: {org_id}, Comms Hub Responded HTTP {response.status_code}")
            continue

        db.updateBlackFlagNotified(conn, org_id)

    conn.close()
    logger.info("Finished safety check.")

def service_loop():
    logger.info("Start service loop")

    while True:
        safety_check()
        validate()
        time.sleep(60)


"""
Classes that provide files and classes to the ImgDataset.
"""
__authors__ = "Nicu Tofan"
__copyright__ = "Copyright 2015, Nicu Tofan"
__credits__ = ["Nicu Tofan"]
__license__ = "3-clause BSD"
__maintainer__ = "Nicu Tofan"
__email__ = "nicu.tofan@gmail.com"

import csv
import functools
import Image
import logging
import numpy
import os
from pylearn2.utils import as_floatX
from pylearn2.utils.string_utils import preprocess
from pylearn2.utils.rng import make_np_rng

class Provider(object):
    """
    Abstract class providing the interface for data providers.

    If implementations below are not suited for your needs implement a
    subclass of this class. The methods to be implemented are ``next()``
    and ``category()``.

    Parameters
    ----------
    categories : list, optional
        A list of categories; if empty the provider will use `get_categs()`.
    """
    def __init__(self, categs=None):
        if not categs is None:
            if isinstance(categs, Provider):
                self.categs = categs.categories()
            else:
                self.categs = categs
        super(Provider, self).__init__()

    def next(self):
        """
        Either gets the name of next file or throws StopIteration.

        Returns
        -------
        fpath : str
            Path for next file to be considered.
        """
        raise NotImplementedError()

    def __next__(self):
        """
        Get next file.

        Returns
        -------
        fpath : str
            Path for next file to be considered.
        """
        return self.next()

    def cnext(self):
        """
        Gets next file; if the end is reached first file is returned
        """
        raise NotImplementedError()

    def __len__(self):
        """
        Number of files stored inside.

        Returns
        -------
        len : int
            A non-negative integer.
        """
        return len(self.everything())

    def __iter__(self):
        """
        Standard interface for iterators.
        """
        return self

    def setup(self, dataset):
        """
        Called by the dataset once it initialized itself.
        """
        pass
        #assert isinstance(dataset, ImgDataset)

    def tear_down(self):
        """
        Called by the dataset fromits tear_down() method.
        """
        pass

    @staticmethod
    def normalize_image(img):
        """
        Converts an image into the format used by ``read()``.
        """
        if img.mode == 'LAB' or img.mode == 'HSV':
            raise ValueError('%s image mode is not supported' % img.mode)
        img = img.convert('RGBA')
        imarray = as_floatX(numpy.array(img)) / 255.0
        assert numpy.all(imarray >= 0.0) and numpy.all(imarray <= 1.0)
        assert len(imarray.shape) == 3
        assert imarray.shape[2] == 4
        return imarray

    def read(self, f_path):
        """
        Read a file and return an image and its category as a tuple.

        This method only reads internal images. For a more general method
        see `read_image()`.

        The returned array has (0, 1, 'c') shape, with c always 4
        (r, g, b, a) and datatype being floatX.

        Parameters
        ----------
        f_path : str
            The path as returned by a previous ``next()`` call.

        Returns
        -------
        result : tuple
            First member is the image as an numpy array,
            second is the category.
        """
        categ = self.category(f_path)
        logging.debug('generator reading [%s] from %s', categ, f_path)
        img = Image.open(f_path)
        return Provider.normalize_image(img), categ

    def read_image(self, image, internal=True):
        """
        Read an image and convert it in the format used by Provider.

        Parameters
        ----------
        image : str or Image.Image or numpy.ndarray
            The path for the image or the image itself.
        internal : bool
            Interpretation of the ``image`` parameter if it is a string:
            - If ``internal`` is True the image is read
              using `Provider.read()` (it is  assumed it is
              one of dataset's own image).
            - If False it is interpreted as being a
              path and is read using Image.

        The returned array is normalized to (0, 1, 'c') shape, with c always 4
        (r, g, b, a) and datatype being floatX.

        Returns
        -------
        result : numpy.ndarray
            The resulted image.
        """
        if isinstance(image, basestring):
            if internal:
                image = self.read(image)[0]
            else:
                image = Image.open(image)
                image = Provider.normalize_image(image)
        elif isinstance(image, numpy.ndarray):
            assert len(image.shape) == 3
            image = as_floatX(image)
        else:
            assert isinstance(image, Image.Image)
            image = Provider.normalize_image(image)
        # image is an array
        if image.shape[2] == 3:
            app = numpy.ones(shape=(image.shape[0], image.shape[1], 1),
                             dtype=image.dtype) * 255
            numpy.append(image, app, axis=2)
        assert numpy.all(image >= 0.0) and numpy.all(image <= 1.0)
        return image

    def category(self, f_path):
        """
        Given a file name returned by next finds the category for that file.

        Parameters
        ----------
        iterator_cls : class
            An iterator class that inherits from SubsetIterator

        Returns
        -------
        categ : str
            A string indicating the category.
        """
        raise NotImplementedError()

    def categ_len(self):
        """
        Number of categories for this Provider.
        """
        return len(self.categories())

    def categories(self):
        """
        Get the list of categories.

        The list should always be the same for the lifetime of the instance
        as the indices are used to identify the class.

        Returns
        -------
        categs : list of str
            A string indicating the category.
        """
        if hasattr(self, 'categs'):
            return self.categs
        else:
            return self.get_categs()

    def get_categs(self):
        """
        Get a list of all categories in this provider.

        The implementation in the base class caches the result so, in you
        later want to insert images make sure to update the category
        yourself.
        """
        self.categs = []
        for key in self:
            ctg = self.category(key)
            if not ctg in self.categs:
                logging.warning('category %d: %s', len(self.categs), str(ctg))
                self.categs.append(ctg)

        return self.categs

    def categ2int(self, categ):
        """
        Converts a string representing a category name into an integer.

        Parameters
        ----------
        categ : str
            A category name.

        Returns
        -------
        idx : int
            The index of this category in the list of categories.
        """
        idx = self.categories().index(categ)
        if idx == -1:
            raise ValueError('%s is not among the '
                             'categories of the Provider' % categ)
        return idx

    def int2categ(self, idx):
        """
        Converts a string representing a category name into an integer.

        Parameters
        ----------
        idx : int
            The numeric id of the catgory.

        Returns
        -------
        categ : str
            The category coresponding to provided integer..
        """

        ctgs = self.categories()
        if idx < 0 or idx >= len(ctgs):
            raise ValueError('%d is outside the valid range '
                             'for categories of the Provider' % idx)
        return ctgs[idx]

    def everything(self):
        """
        The whole set of files and categories as a dictionary.

        Returns
        -------
        retdict : dict
            A dictionary where keys are files and values are categories.
        """
        retdict = {}
        for key in self:
            retdict[key] = self.category(key)
        return retdict

    def __hash__(self):
        """
        Called by built-in function hash() and for operations on members
        of hashed collections including set, frozenset, and dict.
        """
        hdict = hash(self.__class__.__name__)
        evr = self.everything()
        for k in evr:
            hdict = hdict ^ (hash(k) << 1) ^ (hash(evr[k]) << 2)
        return  hdict

    def get(self, offset, count):
        """
        Get the files in a given range.

        The method does not throw errors if the range is outside allowed
        range. instead, the list of files is treated as a ring.

        Parameters
        ----------
        offset : int
            Zero-based indes of the first file to retreive.
        count : int
            Number of files to retreive.
        """
        assert offset >= 0
        assert count > 0
        evrt = self.everything()
        evrt_len = len(evrt)
        keys = evrt.keys()

        result = []
        while count > 0:
            # bring offset in valid range
            offset = offset % evrt_len
            # max number of items
            valid_range = evrt_len - offset
            # how many are we going to get on this round
            this_round = min(valid_range, count)
            result.extend(keys[offset:offset+this_round])

            offset = offset + this_round
            count = count - this_round

        return result

class DictProvider(Provider):
    """
    A provider based on a dictionary.

    Parameters
    ----------
    data : dict of strings
        A dictionary with the keys being file paths and values being
        the category.
    categs : list, optional
        A list of categories; if empty the provider will use `get_categs()`.
    """
    def __init__(self, data, categs=None):
        super(DictProvider, self).__init__(categs)

        #: The dictionary that maps paths to classes.
        self.data = data

        #: The list of paths used for iteration purposes.
        self.keys_iter = data.keys().__iter__()

    @functools.wraps(Provider.__iter__)
    def __iter__(self):
        return self.data.keys().__iter__()

    @functools.wraps(Provider.__iter__)
    def __len__(self):
        return len(self.data)

    @functools.wraps(Provider.next)
    def next(self):
        return self.keys_iter.next()

    @functools.wraps(Provider.category)
    def category(self, f_path):
        return self.data[f_path]

    @functools.wraps(Provider.everything)
    def everything(self):
        return self.data

    @functools.wraps(Provider.cnext)
    def cnext(self):
        try:
            return self.next()
        except StopIteration:
            self.keys_iter = self.data.keys().__iter__()
            return self.next()

    def __getstate__(self):
        """
        Help pickle this instance.
        """
        if not hasattr(self, 'categs'):
            self.get_categs()
            assert hasattr(self, 'categs')

        return {'data': self.data, 'categs': self.categs}

    def __setstate__(self, state):
        """
        Help un-pickle this instance.
        """
        self.data = state['data']
        self.categs = state['categs']

        #: The list of paths used for iteration purposes.
        self.keys_iter = self.data.keys().__iter__()

class CsvProvider(DictProvider):
    """
    A provider based on a comma separated values file.

    The constructor will read the file and construct a dictionary that is
    later used by the DictProvider superclass.

    Parameters
    ----------
    csv_path : str
        The path where the csv file can be found.
    has_header : bool
        Skip first row if true. ``has_header`` must be True if ``col_path``
        or ``col_class`` are strings.
    col_path : str or int, optional
        If this is an integer, it represents the zero based index of the column
        that contains image paths.
        If this is a string, it represents the name of the column; the name
        will be searched in the header (so it is asserted the ``has_header``
        is true).
    col_class : str or int, optional
        If this is an integer, it represents the zero based index of the column
        that contains image paths.
        If this is a string, it represents the name of the column; the name
        will be searched in the header (so it is asserted the ``has_header``
        is true).
    delimiter : str, optional
        A one-character string used to separate fields. It defaults to ``,``.
    quotechar : str, optional
        A one-character string used to quote fields containing special
        characters, such as the delimiter or quotechar,
        or which contain new-line characters. It defaults to ``"``.
    skip_first : int, optional
        Skip this many rows at the beginning of the file; this does
        not include the header row that is managed separatelly.
    skip_last : int, optional
        Skip this many rows at the end of the file.
    categs : list, optional
        A list of categories; if empty the provider will use `get_categs()`.

    Notes
    -----
    The class is only interested in two columns: the class and the path.
    Both results are passed through pylearn2.utils.string_utils.preprocess
    so you can use environment variables to customize their final value.

    The paths inside csv file, if relative, are considered relative to
    the path of the .csv file.

    For the last entries to be removed the entries are parsed
    and added to a list. They are later removed from the result.
    This is in order to avoid reading the csv file twice but
    it has the downside that the entries to be skipped stil
    have to be valid (enough columns).
    """
    def __init__(self, csv_path, col_path=1, col_class=0, has_header=False,
                 delimiter=',', quotechar='"',
                 skip_first=0, skip_last=0, categs=None):

        assert skip_first >= 0
        assert skip_last >= 0

        if (isinstance(col_class, basestring) or
                isinstance(col_path, basestring)):
            # we will need to locate the indices so we need a header
            assert has_header
        else:
            col_min = max(col_path, col_class)

        # make the path absolute and extract base directory
        csv_path = os.path.abspath(csv_path)
        csv_base = os.path.split(csv_path)[0]

        # collect data in a dictionary
        data = {}
        file_list = []
        with open(csv_path, 'rt') as fhand:
            csvr = csv.reader(fhand,
                              delimiter=delimiter,
                              quotechar=quotechar)

            # collect all rows
            for i, row in enumerate(csvr):
                if len(row) == 0:
                    continue
                if has_header:
                    # find the actual indices for columns of interest
                    if isinstance(col_class, basestring):
                        col_class = row.index(col_class)
                    if isinstance(col_path, basestring):
                        col_path = row.index(col_path)
                    has_header = False
                    col_min = max(col_path, col_class)
                    skip_first = skip_first + 1
                    continue

                if i < skip_first:
                    continue

                if len(row) <= col_min:
                    err = '%s[%d] should have at least %d ' + \
                          'columns but it only has %d'
                    raise ValueError(err % (csv_path, i, col_min, len(row)))
                else:
                    class_name = preprocess(row[col_class]).strip().lower()
                    fpath = preprocess(row[col_path])
                    if len(fpath) > 0:
                        if not os.path.isabs(fpath):
                            fpath = os.path.join(csv_base, fpath)
                    data[fpath] = class_name
                    if skip_last > 0:
                        file_list.append(fpath)

        # remove last entries
        if skip_last > 0:
            if len(file_list) <= skip_last:
                data = {}
            else:
                for f2rem in file_list[-skip_last:]:
                    data.pop(f2rem, None)

        # everything else is provided by DictProvider
        super(CsvProvider, self).__init__(data, categs)

class RandomProvider(DictProvider):
    """
    A provider that has 100 random images labelled 1 to 100

    The constructor will create the random images and will construct
    a dictionary that is used by the DictProvider superclass.

    Parameters
    ----------
    rng : `numpy.random.RandomState` or seed, optional
        Seed for random number generator or an actual RNG.
    """
    def __init__(self, rng=None, count=100, alpha=False, size=(128, 128)):
        self.rng = make_np_rng(rng)
        self.content = {}
        self.rng = rng
        self.count = count
        self.alpha = alpha
        self.size = size
        data = self.prepare()

        # everything else is provided by DictProvider
        super(RandomProvider, self).__init__(data)

    @functools.wraps(Provider.__iter__)
    def __len__(self):
        return self.count

    @functools.wraps(Provider.read)
    def read(self, f_path):
        logging.debug('generator reading file %s', f_path)
        categ = f_path
        img = self.content[f_path]
        return Provider.normalize_image(img), categ

    def __getstate__(self):
        """
        Help pickle this instance.
        """
        state = {}
        state['count'] = self.count
        state['alpha'] = self.alpha
        state['size'] = self.size
        return state

    def __setstate__(self, state):
        """
        Help un-pickle this instance.
        """
        self.count = state['count']
        self.alpha = state['alpha']
        self.size = state['size']
        self.data = self.prepare()
        self.keys_iter = self.data.keys().__iter__()

    def prepare(self):
        """
        Initializes the instance according to internal attributes.
        """
        self.content = {}
        data = {}
        modes = ['RGBA', 'RGB', '1', 'L', 'P', 'CMYK', 'I', 'F']
        channels = 4 if self.alpha else 3
        for i in range(self.count):
            file_name = str(i+1)
            imarray = numpy.random.rand(self.size[0],
                                        self.size[1],
                                        channels) * 255
            img = Image.fromarray(imarray.astype('uint8'))
            self.content[file_name] = img.convert(modes[i % len(modes)])
            data[file_name] = file_name
        return data

class DeDeMaProvider(Provider):
    """
    A provider based on a DenseDesignMatrix.

    Parameters
    ----------
    data : dict of strings
        A dictionary with the keys being file paths and values being
        the category.
    """
    def __init__(self, dedema):
        self.dedema = dedema
        self.dataset_iter = range(dedema.get_num_examples()).__iter__()
        super(DeDeMaProvider, self).__init__()

    @functools.wraps(Provider.__iter__)
    def __iter__(self):
        return self.dataset_iter

    @functools.wraps(Provider.__iter__)
    def __len__(self):
        return self.dedema.get_num_examples()

    @functools.wraps(Provider.next)
    def next(self):
        return self.dataset_iter.next()

    @functools.wraps(Provider.category)
    def category(self, f_path):
        f_path = int(f_path)
        assert not self.dedema.y is None
        return self.dedema.y[f_path]

    @functools.wraps(Provider.everything)
    def everything(self):
        rang = range(self.dedema.get_num_examples())
        data = {}
        for i in rang:
            data[i] = self.dedema.y[i]
        return data

    @functools.wraps(Provider.cnext)
    def cnext(self):
        try:
            return self.next()
        except StopIteration:
            rang = range(self.dedema.get_num_examples())
            self.dataset_iter = rang.__iter__()
            return self.next()

    def __getstate__(self):
        """
        Help pickle this instance.
        """
        return {'data': self.dedema}

    def __setstate__(self, state):
        """
        Help un-pickle this instance.
        """
        self.dedema = state['data']
        self.dataset_iter = self.dedema.iterator()

    @functools.wraps(Provider.read)
    def read(self, f_path):
        f_path = int(f_path)
        categ = self.category(f_path)
        ddm = self.dedema.get_design_matrix()
        view_conv = self.dedema.view_converter.design_mat_to_topo_view
        exm = view_conv(ddm[f_path].reshape(1, ddm.shape[1]))
        return exm[0], categ


from exceptions import *
from dakota_file import DakotaFile
import numpy as np
import xarray as xr
import main
import os

my_netcdf = DakotaFile()

name = 'test_scan1'
var_type = 'scan'

# Make a dictionary of the needed data
lower_bounds  = np.array( [ 1.0, 3.0 ] )
upper_bounds  = np.array( [ 2.0, 4.0 ] )
partitions    = np.array( [ 1,   1   ] )

variable_dict = { 'lower_bounds' : lower_bounds, 'upper_bounds' : upper_bounds, 'partitions' : partitions }

my_netcdf.add_variable_from_dict( name, var_type, variable_dict )

name = 'test_scan2'
var_type = 'scan'

lower_bounds = 0.0
upper_bounds = 1.0
partitions   = 1

dataset = xr.Dataset( {'lower_bounds' : lower_bounds, 'upper_bounds' : upper_bounds, 'partitions' : partitions } )

my_netcdf.add_variable_from_dataset( name, var_type, dataset )

name = 'test_scan3'
var_type = 'scan'

variable_dict = { 'lower_bounds' : np.array(10), 'upper_bounds' : np.array(15), 'partitions' : np.array(1) }

my_netcdf.add_variable_from_dict( name, var_type, variable_dict )

filename = 'DAKOTA.nc'
my_netcdf.write(filename)


import sys
import atexit
import logging

from typing import Optional

from psycopg import connect, Connection

from .base import HandlerBase

logger = logging.getLogger()

class PostgresHandler(HandlerBase):
    def __init__(self, dsn: Optional[str] = None, connection: Optional[Connection] = None, setup_script: Optional[str] = None, BATCH_SIZE: int = 1000):
        super().__init__()
        self.conn: Connection = self._setup_connection(dsn=dsn, connection=connection)
        self.ops: list[tuple] = []
        self.BATCH_SIZE: int = BATCH_SIZE
        self._setup_db(setup_script=setup_script)

        if connection is None:
            # Test database will close cleanly causing this to fail, so only register the exit handler if we're not using a test database.
            atexit.register(self._close)

    def _close(self):
        """
        Run any outstanding operations in our queue and then close the connection to the database.
        """
        self._run_ops(self.ops)
        self.conn.close()

    def _setup_connection(self, dsn: Optional[str] = None, connection: Optional[Connection] = None) -> Connection:
        """
        Will return the provided connection if one exists (useful for testing). Otherwise will attempt to connect to the provided DSN.

        Parameters
        ----------
        dsn: str
            Postgresql connection string.
        connection: Connection
            Existing psycopg3 database connection.

        Returns
        -------
        Connection
            psycopg3 database connection.
        """
        try:
            logger.info("Connecting to database.")

            if connection is not None:
                if not isinstance(connection, Connection):
                    raise Exception("Provided connection is not a psycopg3 connection.")

                return connection

            return connect(dsn, autocommit=False)
        except Exception as e:
            logger.error("Could not connect to database.")
            logger.error(e)
            sys.exit(1)

    def _setup_db(self, setup_script: Optional[str] = None):
        """
        Run a database setup script.

        Parameters
        ----------
        setup_script: str
            Path to a SQL file which will be executed by the database.
        """
        if setup_script is not None:
            try:
                logger.info("Running db setup script.")

                with open(setup_script, 'r') as f:
                    self.queue_op(f.read(), run_now=True)
            except (IOError, FileNotFoundError) as e:
                logger.error("Could not open the database setup script.")

    def queue_op(self, sql: str, args: Optional[tuple] = None, run_now: bool = False) -> None:
        """
        Adds the provided sql and args as a tuple to an internal ops list. When the length of this list is BATCH_SIZE, execute them all.

        Parameters
        ----------
        sql: str
            SQL to be executed
        args: tuple
            Args tuple which will be injected into sql
        run_now: bool
            Whether to run the current operation now or queue it up for later.
        """
        self.ops.append((sql, args))

        # By default, execute commands in batches of BATCH_SIZE, or override with run_now
        if len(self.ops) == self.BATCH_SIZE or run_now:
            self._run_ops(self.ops)
            self.ops = []

    def _run_ops(self, ops: list[tuple]):
        """
        Opens a transaction, executes all of the (sql, params) tuples in ops inside of the transaction, and commits it.

        ops: list[tuple]
            A list of (sql, args) pairs to be executed.
        """
        try:
            with self.conn.transaction():
                with self.conn.cursor() as cur:
                    logger.info(f"Running batch of {len(ops)}.")

                    for (sql, args) in ops:
                        if args is not None:
                            cur.execute(sql, args)
                        else:
                            cur.execute(sql)
        except Exception as e:
            logger.error("Could not execute database operation.")
            raise e


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Feb 13 12:35:03 2019

@author: calbert
"""
import time
import numpy as np
import matplotlib.pyplot as plt
from fffi import fortran_library, fortran_module

pi = 3.14159265358979
c = 2.9979e10
e_charge = 4.8032e-10
e_mass = 9.1094e-28
p_mass = 1.6726e-24
ev = 1.6022e-12

bmod_ref = 5.0e4
Z_charge = 2
m_mass = 4
E_kin = 3.5e6
trace_time = 1.0e-2
# rbig=rlarm*bmod00
# dphi=2.d0*pi/(L1i*npoiper)

v0 = np.sqrt(2.0*E_kin*ev/(m_mass*p_mass))    # alpha particle velocity, cm/s
rlarm = v0*m_mass*p_mass*c/(Z_charge*e_charge*bmod_ref)  # reference gyroradius
tau = trace_time*v0

libneo_orb = fortran_library('simple')

neo_orb = fortran_module(libneo_orb, 'neo_orb_global')
neo_orb.fdef("""
  subroutine init_field(ans_s, ans_tp, amultharm, aintegmode)
    integer :: ans_s, ans_tp, amultharm, aintegmode
  end

  subroutine init_params(Z_charge, m_mass, E_kin, adtau, adtaumax, arelerr)
    integer :: Z_charge, m_mass
    double precision :: E_kin, adtau, adtaumax
    double precision :: arelerr
  end

  subroutine init_integrator(z0)
    double precision, dimension(:), intent(in) :: z0
  end

  subroutine timestep_z(z, ierr)
    double precision, dimension(:) :: z
    integer :: ierr
  end

  subroutine timestep_sympl_z(z, ierr)
    double precision, dimension(:) :: z
    integer :: ierr
  end
""")

new_vmec_stuff = fortran_module(libneo_orb, 'new_vmec_stuff_mod')
new_vmec_stuff.fdef("""
    double precision :: rmajor, h_theta, h_phi
""")

cut_detector = fortran_module(libneo_orb, 'cut_detector_global')
cut_detector.fdef("""
    subroutine init(z)
      double precision, dimension(:), intent(in) :: z
    end

    subroutine trace_to_cut(z, var_cut, cut_type, ierr)
      double precision, dimension(:), intent(inout) :: z
      double precision, dimension(:), intent(inout) :: var_cut
      integer, intent(out) :: cut_type
      integer, intent(out) :: ierr
    end
""")

libneo_orb.compile()
neo_orb.load()
new_vmec_stuff.load()

# %%
neo_orb.init_field(5, 5, 3, 1)

# %%
npoiper2 = 24                       # interation points per field period
rbig = new_vmec_stuff.rmajor*1.0e2  # major radius in cm
dtaumax = 2.0*pi*rbig/npoiper2      # maximum time step for integrator
dtau = dtaumax                   # time step for output
neo_orb.init_params(Z_charge, m_mass, E_kin, dtau, dtaumax, 1e-8)
# %%
cut_detector.load()

# %%
ntimstep = 10000
neo_orb.firstrun = 1  # to re-initialize symplectic integrator in case

s = 0.5
th = 0.0
ph = 0.314
lam = 0.22

z = np.array([s, th, ph, 1.0, lam])
neo_orb.init_integrator(z)

zs = np.empty([4, ntimstep+1])
zs[:, 0] = z[[0, 1, 2, 4]]

t = time.time()
ierr = 0  # doesn't work yet with pass-by-reference
for kt in range(ntimstep):
    neo_orb.timestep_sympl_z(z, ierr)
    zs[:, kt+1] = z[[0, 1, 2, 4]]
print(z)
print('time elapsed: {} s'.format(time.time() - t))

cut_detector.init(z)
ierr = 0
var_cut = np.zeros(6)
ncut = 1000
# variables to evaluate at tip: z(1..5), par_inv
var_tips = np.empty([6, ncut])

t = time.time()
for k in np.arange(ncut):
    cut_type = 0
    cut_detector.trace_to_cut(z, var_cut, cut_type, ierr)
    var_tips[:, k] = var_cut
print(z)
print('time elapsed: {} s'.format(time.time() - t))

plt.plot(zs[0, :]*np.cos(zs[1, :]), zs[0, :]*np.sin(zs[1, :]))
plt.plot(var_tips[0, :]*np.cos(var_tips[1, :]),
         var_tips[0, :]*np.sin(var_tips[1, :]), ',')


import json

def stringify(data, depth: int):
    """
    Convert data to strings with correct indent.
    """
    if not isinstance(data, dict):

        return save_bool_format(data)

    int_indent = build_indent(depth + 1)
    ext_indent = build_indent(depth)
    strings = ['{']

    for key, val in data.items():
        strings.append(f"\n{int_indent}  {key}: "
                       f"{stringify(val, depth + 1)}"
                       )
    strings.append(f"\n  {ext_indent}}}")

    return ''.join(strings)

def build_indent(depth):
    """
    Gen correct indent depend on current node`s depth.
    """
    return ' ' * (depth * 4 - 2)

def format_(tree, depth=0):
    """
    Convert tree with differences to str in stylish format and correct indent.
    Args:
        tree: list of objects with differences,
        depth: current node`s depth (0 by default).
    Returns:
        output: Str in stylish format or exception if diff_tree is incorrect.
    """
    STAT = {'old_value': '- ',
            'new_value': '+ ',
            'removed': '- ',
            'added': '+ ',
            'unchanged': '  ',
            'look_inside': '  '
            }
    indent = build_indent(depth)

    if tree['type'] == 'diff':
        lines = map(lambda child: format_(child, depth + 1),
                    tree.get('children')
                    )
        result = '\n'.join(lines)
        return f'{{\n{result}\n}}'

    elif tree['type'] == 'look_inside':
        lines = map(lambda child: format_(child, depth + 1),
                    tree.get('children')
                    )
        result = '\n'.join(lines)
        output = (f"{indent}{STAT[tree['type']]}{tree['name']}:"
                  f" {{\n{result}\n{indent}  }}"
                  )

        return output

    elif tree['type'] in ('added', 'removed', 'unchanged'):
        output = (f"{indent}{STAT[tree['type']]}{tree['name']}:"
                  f" {stringify(tree.get('value'), depth)}"
                  )

        return output

    elif tree['type'] == 'changed':
        output = (f"{indent}{STAT['old_value']}{tree['name']}:"
                  f" {stringify(tree.get('old_value'), depth)}\n"
                  f"{indent}{STAT['new_value']}{tree['name']}:"
                  f" {stringify(tree.get('new_value'), depth)}"
                  )

        return output

    raise Exception('Ivalid node`s type.')

def save_bool_format(data):
    """
    Change bool-types values to it source (json/yaml) format.

    Args:
        file: data in python format.
    Returns:
           file: data in json/yaml format.
    """
    CHANGE_COLLECTION = {False, True, None}

    if data in CHANGE_COLLECTION:

        return json.dumps(data)

    return data


import torch
from torch import device, cuda
import numpy as np
device = device("cuda" if cuda.is_available() else "cpu")

class State(object):

    def __init__(self, env, region_nr):
        """"
        New state space at beginning of each episode.
        Initializes a state for the given zipcode in the specified region.
        All ambulances are initially available and no accidents have occured yet.
        """
        # Doesn't it initialize a state for the region and not the specific zipcode?
        self.env = env
        self.K = 4
        self.N = len(env.postcode_dic[region_nr])
        self.ambulance_return = {}  # Dictionary with key: when will an ambulance return and value: zip code of base
        self.region_nr = region_nr
        self.waiting_list = []
        self.accident_location_index = -1

        # parameters initialized to pass to NN
        self.bool_accident = [0] * self.N
        # is_base: boolean list of whether an env.postcode_dic[region_nr][i] zip_code is a base
        # nr_ambulances: int list of how many ambulances are available per zip_code
        self.is_base, self.nr_ambulances = self.check_isBase()
        self.travel_time = [0] * self.N  # time from base to accident
        self.delta = env.prob_acc[region_nr]
        self.time = [0] * self.N

        # index of bases (should not be masked)
        self.indexNotMasked = np.array([i for i, e in enumerate(self.is_base) if e == 1])

    def check_isBase(self):
        """
        Find all bases in a region
        :param env:
        :param region_nr:
        :return: boolean list indicating if zip-code has a base or not
        """
        isBase = []
        nr_ambulances = []
        for zip_code in self.env.postcode_dic[self.region_nr]:
            if zip_code in self.env.bases[self.region_nr]:
                isBase.append(1)
                nr_ambulances.append(self.env.bases[self.region_nr][zip_code])
            else:
                isBase.append(0)
                nr_ambulances.append(0)
        return isBase, nr_ambulances

    def update_travel_time(self, accident_loc):
        """
        Takes the accident location zip code and updates the travel_time list.
        :param accident_loc: zip code that the ambulance is sent to
        """
        for i, zip_code in enumerate(self.env.postcode_dic[self.region_nr]):
            if zip_code in self.env.bases[self.region_nr]:
                if self.nr_ambulances[i] > 0:
                    self.travel_time[i] = self.env.distance_time(self.region_nr, zip_code, accident_loc)
                else:
                    self.travel_time[i] = 99999
            else:
                self.travel_time[i] = 99999

    def update_travel_time_greedy(self, accident_loc):
        """
        Takes the accident location zip code and updates the travel_time list.
        :param accident_loc: zip code that the ambulance is sent to
        """
        for i, zip_code in enumerate(self.env.postcode_dic[self.region_nr]):
            if zip_code in self.env.bases[self.region_nr]:
                if self.nr_ambulances[i] > 0:
                    self.travel_time[i] = self.env.distance_time(self.region_nr, zip_code, accident_loc)
                else:
                    self.travel_time[i] = 99999
            else:
                self.travel_time[i] = 99999

    def process_action(self, action, time):
        """
        Takes an action (ambulance sent out) and returns the new state and reward.
        :param action: index of zip code that send out ambulance
        :param time: time of the day in seconds that ambulance was sent out
        :param accident_loc: location of the accident to calculate when an amublance will arrive
        :return reward: minus time from ambulance to the accident
        """
        if action == 0:

            return -100
        else:
            accident_loc = self.get_accident_location()
            self.nr_ambulances[action] -= 1
            if self.nr_ambulances[action] == 0:
                self.indexNotMasked = self.indexNotMasked[self.indexNotMasked != action]
            total_travel_time = self.env.calculate_ttt(self.region_nr, self.env.postcode_dic[self.region_nr][action],
                                                       accident_loc)
            self.ambulance_return.update({total_travel_time + time: action})
        if self.travel_time[action] == 0:
            return 10000 / 100
        return 10000/self.travel_time[action]
        #return -self.travel_time[action]

    def process_action_greedy(self, action, time):
        """
        Takes an action (ambulance sent out) and returns the new state and reward.
        :param action: index of zip code that send out ambulance
        :param time: time of the day in seconds that ambulance was sent out
        :param accident_loc: location of the accident to calculate when an amublance will arrive
        :return reward: minus time from ambulance to the accident
        """
        if action == 0 or self.nr_ambulances[action] == 0:

            return -100
        else:
            accident_loc = self.get_accident_location()
            self.nr_ambulances[action] -= 1
            if self.nr_ambulances[action] == 0:
                self.indexNotMasked = self.indexNotMasked[self.indexNotMasked != action]
            total_travel_time = self.env.calculate_ttt(self.region_nr, self.env.postcode_dic[self.region_nr][action],
                                                       accident_loc)
            self.ambulance_return.update({total_travel_time + time: action})
        if self.travel_time[action] == 0:
            return 10000 / 100

        return 10000/self.travel_time[action]
        #return -self.travel_time[action]

    def update_state(self, time, zip_code_index):
        if self.accident_location_index != -1:
            self.bool_accident[self.accident_location_index] = 0
        self.bool_accident[zip_code_index] = 1
        self.accident_location_index = zip_code_index
        for i in range(self.N):
            self.time[i] = time

        self.update_travel_time(self.get_accident_location())

    def update_state_greedy(self, time, zip_code_index):
        if self.accident_location_index != -1:
            self.bool_accident[self.accident_location_index] = 0
        self.bool_accident[zip_code_index] = 1
        self.accident_location_index = zip_code_index
        for i in range(self.N):
            self.time[i] = time

        self.update_travel_time_greedy(self.get_accident_location())

    def get_accident_location(self):
        return self.env.postcode_dic[self.region_nr][self.accident_location_index]

    def get_torch(self):
        """
        Transform state object into a KxN torch, where K = number of parameters and N = number of zipcodes
        :return:
        """
        return torch.transpose(torch.tensor([self.nr_ambulances,
                                             self.travel_time,
                                             self.time,
                                             self.is_base], device=device), 0, 1)

    def __deepcopy__(self, memodict={}):
        copy_object = State(self.env, self.region_nr)

        copy_object.ambulance_return = self.ambulance_return.copy()
        copy_object.waiting_list = self.waiting_list.copy()

        # parameters initialized to pass to NN
        copy_object.bool_accident = self.bool_accident.copy()
        # is_base: boolean list of whether an env.postcode_dic[region_nr][i] zip_code is a base
        # nr_ambulances: int list of how many ambulances are available per zip_code
        copy_object.is_base, copy_object.nr_ambulances = self.is_base.copy(), self.nr_ambulances.copy()
        copy_object.travel_time = self.travel_time.copy()
        copy_object.delta = self.delta.copy()
        copy_object.time = self.time.copy()

        # index of bases (should not be masked)
        copy_object.indexNotMasked = self.indexNotMasked.copy()
        return copy_object