#!/usr/bin/env python3
# Christopher Vollmers
# Roger Volden

import sys
import numpy as np

content_file = sys.argv[1]
out_path = sys.argv[2]
cutoff = float(sys.argv[3])
genome_file = sys.argv[4]
refine = sys.argv[5]

minimum_read_count = 3
minimum_read_coverage = 2
splice_site_width = 5

def scan_for_best_bin(entry, distance_range, iterator_shift, density_dict,
                      base_cutoff_min, base_cutoff_max,
                      peak_areas, chromosome, side):
    '''
    Inputs:
        entry (dict): bound:[features] for whichever chromosome
        distance_range (list): double the splice site width
        iterator_shift (int): determines which way you iterate (f/r)
        density_dict (dict): same as histo_left/right_bases but for a chromosome
        base_cutoff_min (int, 0): arbitrary lower cutoff
        base_cutoff_max (int, 5): arbitrary upper cutoff
        peak_areas (dict): keeps track of chromosome regions with peaks
        chromosome (str): which chromosome
        side (str): which direction to look (l/r)
    Outputs:
        best_extra_list (list): should be int, counts how many features
        peak_center (int): position of the center of the peak
        bases (list): distance b/w aligned blocks
        coverage_area (list): contains positions that were covered
        best_direction_l/r (dict): keeps track of highest counts in dist range
    '''

    best_extra_list, peak_center, bases = [], 0, []
    coverage_area, best_direction_l, best_direction_r = [], [], []
    for x in distance_range:
        extra_list_bases, extra_list_expression = [], []
        direction_l, direction_r = {}, {}
        coverage_set = []
        called = False
        for y in distance_range:
            try:
               called = peak_areas[chromosome][side][entry+x+y]
            except KeyError:
               pass
        if not called:
             highest_y, highest_y_pos = 0, 0
             for y in distance_range:
                try:
                    for item in density_dict[entry+x+y]:
                        extra_list_bases.append(item[0])
                        extra_list_expression.append(1)
                        if not direction_l.get(item[4]):
                            direction_l[item[4]] = 1
                        else:
                            direction_l[item[4]] += 1
                        if not direction_r.get(item[5]):
                            direction_r[item[5]] = 1
                        else:
                            direction_r[item[5]] += 1
                        for covered_position in item[3]:
                            coverage_set.append(covered_position)
                except:
                    pass

        if base_cutoff_min <= np.median(extra_list_bases) <= base_cutoff_max:
            if sum(extra_list_expression) > sum(best_extra_list):
                best_extra_list = extra_list_expression
                peak_center = entry + x
                bases = extra_list_bases
                coverage_area = coverage_set
                best_direction_l = direction_l
                best_direction_r = direction_r

    return best_extra_list, peak_center, bases, coverage_area, \
           best_direction_l, best_direction_r

def determine_coverage(coverage_area, chromosome, reverse,
                       peak_center, histo_coverage):
    '''
    Inputs:
        coverage_area (list): contains positions that were covered
        chromosome (str): which chromosome
        reverse (bool): determines which way you read the sequence
        peak_center (int): position of the center of the peak
        histo_coverage (dict): approximate positions of alignment starts
    Outputs:
        coverage (int): how many times the most covered feature was covered
        coverage_area (list): contains positions that were covered
    '''

    coverage = [0]
    coverage_area2 = []
    for covered_position in set(coverage_area):
        if coverage_area.count(covered_position) > 1:
            coverage_area2.append(covered_position)

    coverage_area = sorted(coverage_area2, reverse=reverse)
    counter = 0
    for base_f in coverage_area:
        count = 0
        if not reverse:
            if base_f > peak_center:
                count = 1
        elif reverse:
            if base_f < peak_center:
                count = 1
        if count == 1:
            if counter <= 3:
                counter += 1
                base_f = myround(base_f)
                try:
                    coverage.append(histo_coverage[chromosome][base_f])
                except KeyError:
                    pass
            else:
                break
    coverage = max(coverage)
    return coverage, coverage_area

def read_seq_file(inFile):
    '''Reads in FASTA files, returns a dict of header:sequence'''
    readDict = {}
    for line in open(inFile):
        line = line.rstrip()
        if not line:
            continue
        if line.startswith('>'):
            readDict[line[1:]] = ''
            lastHead = line[1:]
        else:
            readDict[lastHead] += line
    return readDict

def myround(x, base=10):
    '''Rounds to the nearest base'''
    return int(base * round(float(x)/base))

def find_peaks(density_dict, out, peaks, reverse, cutoff, base_cutoff_min,
               base_cutoff_max, histo_coverage, side, peak_areas, chromosome):
    '''
    Inputs:
        density_dict (dict): same as histo_left/right_bases but for a chromosome
        out (file): output bed file
        peaks (int): how many peaks
        reverse (bool): determines which way you read the sequence
        cutoff (float): inversely proportional to coverage
        base_cutoff_min (int, 0): arbitrary lower cutoff
        base_cutoff_max (int, 5): arbitrary upper cutoff
        histo_coverage (dict): approximate positions of alignment starts
        side (str): which direction to look (l/r)
        peak_areas (dict): keeps track of chromosome regions with peaks
        chromosome (str): which chromosome
    Outputs:
        peaks and peak_areas
    '''

    if not reverse:
        distance_range = range(-splice_site_width, splice_site_width)
        iterator_shift = 1
    if reverse:
        distance_range = range(splice_site_width, -splice_site_width, -1)
        iterator_shift =- 1

    entry_list = []
    for entry in density_dict:
      entry_list.append([entry, density_dict[entry]])

    for entry, density in sorted(entry_list,
                                 key=lambda x: sum(np.array(x[1])[:,2]),
                                 reverse=True):
        if len(density) >= minimum_read_count:
            if not peak_areas[chromosome][side].get(entry):

                best_extra_list, peak_center, bases, \
                coverage_area, best_direction_l, best_direction_r \
                = scan_for_best_bin(entry, distance_range, iterator_shift,
                                    density_dict, base_cutoff_min,
                                    base_cutoff_max, peak_areas,
                                    chromosome, side)

                coverage, coverage_area \
                = determine_coverage(coverage_area, chromosome, reverse,
                                     peak_center, histo_coverage)

                if coverage > 0:
                    proportion = round(sum(best_extra_list)/coverage, 3)
                    print(chromosome + '\t' + str(peak_center - 1) + '\t'
                          + str(peak_center + 1) + '\t' + str(proportion))
                    if proportion > cutoff:
                        try:
                            Left_TSS = best_direction_l['TSS']
                        except:
                            Left_TSS = 0
                        try:
                            Left_TES = best_direction_l['TES']
                        except:
                            Left_TES = 0
                        try:
                            Right_TSS = best_direction_r['TSS']
                        except:
                            Right_TSS = 0
                        try:
                            Right_TES = best_direction_r['TES']
                        except:
                            Right_TES = 0
                        Left_to_Right = Left_TSS + Right_TES
                        Right_to_Left = Left_TES + Right_TSS
                        Type = '-'

                        if Left_to_Right < Right_to_Left and reverse:
                            Type = '3'
                        elif Left_to_Right < Right_to_Left and not reverse:
                            Type = '5'
                        elif Left_to_Right > Right_to_Left and reverse:
                            Type = '5'
                        elif Left_to_Right > Right_to_Left and not reverse:
                            Type = '3'

                        if Type != '-':
                            peaks += 1
                            out.write(chromosome + '\t'
                                      + str(peak_center-splice_site_width)
                                      + '\t' + str(peak_center+splice_site_width)
                                      + '\t' + str(Type) + side + str(peaks) + '_'
                                      + str(peak_center-splice_site_width) + '_'
                                      + str(peak_center+splice_site_width) + '_'
                                      + str(proportion) + '\t' + str(peaks) + '\n')
                            for base in range(peak_center - splice_site_width,
                                              peak_center + splice_site_width):
                                peak_areas[chromosome][side][base] = 1
        else:
            break

    return peaks, peak_areas

def collect_reads(content_file):
    '''
    Takes a content file and returns
        histo_left/right_bases (dict): information around upper/lower bounds
        chromosome_list (set): chromosomes with decent alignments
        histo_coverage (dict): approximate positions of alignment starts
    '''

    histo_left_bases, histo_right_bases = {}, {}
    chromosome_list_left, chromosome_list_right = set(), set()
    histo_coverage = {}
    base_cutoff_min, base_cutoff_max = 0, 5

    for line in open(content_file):
        total = 0
        b = line.strip().split('\t')
        infile = b[0]

        length = 0

        for line in open(infile):
            total += 1
            a = line.strip().split('\t')
            chromosome = a[13]

            if not histo_coverage.get(chromosome):
                histo_coverage[chromosome] = {}
            # matches, dir, query seq name
            score, direction, name = int(a[0]), a[8], a[9]
            coverage = int(name.split('_')[3])
            # collect necessary info from psl file
            if coverage >= minimum_read_coverage:
                length = int(name.split('_')[5].split('|')[0])
                begin, span = int(a[15]), int(a[16])
                blocksizes = a[18].split(',')[:-1]
                blockstarts = a[20].split(',')[:-1]
                readstarts = a[19].split(',')[:-1]

                # account for direction in naming
                if direction == '+':
                    start_seq, end_seq = 'S', 'E'
                    left_match, right_match = 'TSS', 'TES'

                else:
                    start_seq, end_seq= 'E', 'S'
                    left_match, right_match = 'TES', 'TSS'

                coverage_set = set()
                previous_blocksize, previous_start = -1, -1
                previous_blockend = np.inf
                intron, indel, indel1 = 0, 0, 0
                low_bounds, up_bounds = [], []
                aligned_bases = 0
                # for each alignment block
                for x in range(0, len(blocksizes)):
                    blockstart = int(blockstarts[x])
                    blocksize = int(blocksizes[x])
                    readstart = int(readstarts[x])
                    aligned_bases += blocksize
                    blockend = blockstart + blocksize
                    # only take alignment blocks >10bp
                    if blocksize > 10:
                        for y in range(0, blocksize, 10):
                            rounded = myround(blockstart + y)
                            coverage_set.add(rounded)
                        for yy in range(y, blocksize):
                            rounded = myround(blockstart + yy)
                            coverage_set.add(rounded)
                        if previous_start == -1:
                            previous_start = blockstart
                        min_length = 10

                        if blockstart - previous_blockend > 20:
                            previous_start = blockstart

                        if blockend - previous_start > min_length:
                            if intron == 1:
                                up_bounds.append([previous_start, indel1, blockend])
                                low_bounds.append([remember_blockend,
                                                   remember_indel1, remember_start])
                                intron = 0
                            else:
                                try:
                                    next_blockstart = int(blockstarts[x+1])
                                    next_blocksize = int(blocksizes[x+1])
                                    next_readstart = int(readstarts[x+1])

                                    insert = next_blockstart - blockend
                                    if insert > 50:
                                        indel1 = next_readstart \
                                                 - (readstart + blocksize)
                                        remember_blockend = blockend
                                        remember_indel1 = indel1
                                        remember_start = previous_start
                                        intron = 1
                                        previous_start = next_blockstart
                                        blockend = next_blockstart
                                except:
                                    pass
                        previous_blockend = blockend

                for rounded in coverage_set:
                    try:
                        histo_coverage[chromosome][rounded] += 1
                    except:
                        histo_coverage[chromosome][rounded] = 1

                # collect chromosome info either to the left or right of a bound
                if aligned_bases/length > 0.70:
                    for low_bound, indel1, blockend in low_bounds:
                        chromosome_list_left.add(chromosome)
                        if not histo_left_bases.get(chromosome):
                            histo_left_bases[chromosome] = {}
                        if not histo_left_bases[chromosome].get(low_bound):
                            histo_left_bases[chromosome][low_bound] = []
                        histo_left_bases[chromosome][low_bound].append([indel1, begin,
                                                                        span, coverage_set,
                                                                        left_match,
                                                                        right_match])
                    for up_bound, indel1, blockend in up_bounds:
                        chromosome_list_right.add(chromosome)
                        if not histo_right_bases.get(chromosome):
                            histo_right_bases[chromosome] = {}
                        if not histo_right_bases[chromosome].get(up_bound):
                            histo_right_bases[chromosome][up_bound] = []
                        histo_right_bases[chromosome][up_bound].append([indel1, begin,
                                                                        span, coverage_set,
                                                                        left_match,
                                                                        right_match])

    chromosome_list = chromosome_list_left & chromosome_list_right
    return histo_left_bases, histo_right_bases, chromosome_list, histo_coverage

def parse_genome(input_file, left_bounds, right_bounds):
    '''
    Takes a GTF file and dictionaries for left and right bounds
    and returns the same left and right bound dictionaries.

    Left/right bounds: dictionary of dictionaries of lists.
    {chromosome:{direction:[start/stop positions]}}
    '''

    # make a dictionary of transcript ID to its
    # chr name, start pos, stop pos, and direction
    # {transcript_ID:[(chr, start, stop, dir), (...) ...], ...}
    gene_dict = {}
    for line in open(input_file):
        a = line.strip().split('\t')
        if len(a) > 7:
            if a[2] == 'exon':
                testKey = a[8].split('; transcript_id "')[1].split('"')[0]
                if not gene_dict.get(testKey):
                    gene_dict[testKey] = []
                # append chr name, start pos, end pos, strand
                gene_dict[testKey].append((a[0], a[3], a[4], a[6]))

    read_list = []
    for transcript_id in gene_dict:
        transcript_data = gene_dict[transcript_id]

        chromosome = transcript_data[0][0]
        if not right_bounds.get(chromosome):
            left_bounds[chromosome], right_bounds[chromosome]= {}, {}
            left_bounds[chromosome]['5'], right_bounds[chromosome]['5'] = [], []
            left_bounds[chromosome]['3'], right_bounds[chromosome]['3'] = [], []

        # first start
        start = sorted(transcript_data, key=lambda x: int(x[1]))[0][1]
        # last end
        end = sorted(transcript_data, key=lambda x: int(x[2]), reverse=True)[0][2]

        # organize data by direction
        # only take the absolute start and absolute end for each transcript
        for entry in transcript_data:
            if entry[1] != start:
                if entry[3] == '+':
                    right_bounds[chromosome]['3'].append(int(entry[1])-1)
                elif entry[3] == '-':
                    right_bounds[chromosome]['5'].append(int(entry[1])-1)
            if entry[2] != end:
                if entry[3] == '+':
                    left_bounds[chromosome]['5'].append(int(entry[2]))
                if entry[3] == '-':
                    left_bounds[chromosome]['3'].append(int(entry[2]))
    return left_bounds, right_bounds

def make_genome_bins(bounds, side, peaks, chromosome, peak_areas):
    '''
    Takes the bounds for a chromosome, which side, the number of peaks,
    which chromosome, and peak areas and returns the number of peaks as
    well as the peak areas.
        bounds (dict): {chromosome:{direction:[start/stop positions]}}
        side (str): left or right (l/r)
        peaks (int): number of peaks
        chromosome (str): which chromosome
        peak_areas (dict): keeps track of chromosome regions with peaks
    '''

    for type1 in ['5', '3']:
        covered = {}
        position_list = sorted(bounds[type1], key=int)
        for index1 in range(0, len(position_list)):
            if not covered.get(index1):
                sub_list=[]
                sub_list.append(position_list[index1])
                for index2 in range(index1, len(position_list)):
                    if position_list[index2] - max(sub_list) <= splice_site_width:
                        sub_list.append(position_list[index2])
                        covered[index2] = 1
                    else:
                        break
                single = 0
                if len(sub_list) > 1:
                    splice_distances = []
                    for splice_pos in range(0, len(sub_list)-1):
                        splice_distances.append(int(sub_list[splice_pos+1])
                                                -int(sub_list[splice_pos]))
                    if min(splice_distances) > 3:
                        for x in range(0,len(sub_list),1):
                            if x != 0:
                                start = int(sub_list[x]
                                            - ((sub_list[x] - sub_list[x-1])/2))
                            else:
                                start = int(sub_list[x]) - 1
                            if x != len(sub_list) - 1:
                                end = int(sub_list[x]
                                          + ((sub_list[x+1] - sub_list[x])/2))
                            else:
                                end = int(sub_list[x]) + 1

                            out.write(chromosome + '\t' + str(start) + '\t'
                                      + str(end) + '\t' + type1 + side
                                      + str(peaks) + '_' + str(start) + '_'
                                      + str(end) + '_A' + '\t' + str(peaks) + '\n')
                            for base in range(start, end):
                                peak_areas[chromosome][side][base] = 1
                            peaks += 1
                    else:
                         single = 1
                else:
                    single = 1
                if single == 1:
                    start = min(sub_list) - splice_site_width
                    end = max(sub_list) + splice_site_width
                    out.write(chromosome + '\t' + str(start) + '\t' + str(end)
                              + '\t' + type1 + side + str(peaks) + '_'
                              + str(start) + '_' + str(end) + '_A'
                              + '\t' + str(peaks) + '\n')
                    for base in range(start-1, end+1):
                        peak_areas[chromosome][side][base] = 1
                    peaks += 1

    return peaks, peak_areas

left_bounds = {}
right_bounds = {}

left_bounds, right_bounds = parse_genome(genome_file, left_bounds, right_bounds)

Left_Peaks = 0
Right_Peaks = 0

histo_left_bases, histo_right_bases, \
chromosome_list, histo_coverage = collect_reads(content_file)
out = open(out_path + '/SS.bed', 'w')

peak_areas = {}
print(chromosome_list)
for chromosome in chromosome_list:
    peak_areas[chromosome] = {}
    peak_areas[chromosome]['l'] = {}
    peak_areas[chromosome]['r'] = {}
    print(chromosome)
    if 'g' in refine:
        Left_Peaks_old = Left_Peaks
        Right_Peaks_old = Right_Peaks
        Left_Peaks, peak_areas = make_genome_bins(left_bounds[chromosome], 'l',
                                                  Left_Peaks, chromosome, peak_areas)
        Right_Peaks, peak_areas = make_genome_bins(right_bounds[chromosome], 'r',
                                                   Right_Peaks, chromosome, peak_areas)
        print('Annotation-Based',
              Left_Peaks - Left_Peaks_old,
              Right_Peaks - Right_Peaks_old)
    Left_Peaks_old = Left_Peaks
    Right_Peaks_old = Right_Peaks
    Left_Peaks, peak_areas = find_peaks(histo_left_bases[chromosome],
                                        out, Left_Peaks, True, cutoff,
                                        0, 5, histo_coverage, 'l',
                                        peak_areas, chromosome)
    Right_Peaks, peak_areas = find_peaks(histo_right_bases[chromosome],
                                         out, Right_Peaks, False, cutoff,
                                         0, 5, histo_coverage, 'r',
                                         peak_areas, chromosome)
    print('Read-Based',
          Left_Peaks - Left_Peaks_old,
          Right_Peaks - Right_Peaks_old)
print(Left_Peaks)
print(Right_Peaks)


# Authors: Stephane Gaiffas <stephane.gaiffas@gmail.com>
# License: BSD 3 clause

"""
This module contains all required container types for nodes (including the
node numpy dtypes and several functions operating locally in nodes.
"""

from math import log
import numpy as np
from numba import from_dtype, jit, boolean, uint8, intp, uintp, float32, void
from numba.experimental import jitclass

from ._utils import get_type, sample_without_replacement
from ._tree_context import TreeClassifierContextType, TreeRegressorContextType

# This data-type describes all the information saved about a node.
# It's used in _tree.py where the Tree dataclass contains an attribute called nodes
# that uses this data type.
node_dtype = np.dtype(
    [
        # Index of the node in the nodes array
        ("node_id", np.uintp),
        #
        # Index of the parent node
        ("parent", np.uintp),
        #
        # Index of the left child node (it's TREE_LEAF is node is a leaf)
        ("left_child", np.intp),
        #
        # Index of the right child node (it's TREE_LEAF is node is a leaf)
        ("right_child", np.intp),
        #
        # Is the node a leaf ?
        ("is_leaf", np.bool),
        #
        # Is the node a left child ? (it's ? for the root node)
        ("is_left", np.bool),
        #
        # Depth of the node in the tree
        ("depth", np.uintp),
        #
        # Feature used for splitting the node (it's ??? if node is a leaf)
        ("feature", np.uintp),
        #
        # Continuous threshold used for splitting the node (not used for now)
        ("threshold", np.float32),
        #
        # Index of the bin threshold used for splitting the node
        ("bin_threshold", np.uint8),
        #
        # Impurity of the node. Used to avoid to split a "pure" node (with impurity=0).
        #   Note that impurity is computed using training samples (as all tree-growing
        #   related things)
        ("impurity", np.float32),
        #
        # Validation loss of the node, computed on validation (out-of-the-bag) samples
        ("loss_valid", np.float32),
        #
        # Logarithm of the subtree aggregation weight
        ("log_weight_tree", np.float32),
        #
        # Number of training samples in the node
        ("n_samples_train", np.uintp),
        #
        # Number of validation (out-of-the-bag) samples in the node
        ("n_samples_valid", np.uintp),
        #
        # Weighted number of training samples in the node
        ("w_samples_train", np.float32),
        #
        # Weighted number of validation (out-of-the-bag) samples in the node
        ("w_samples_valid", np.float32),
        #
        # Index of the first training sample in the node. We have that
        #   partition_train[start_train:end_train] contains the indexes of the node's
        #   training samples
        ("start_train", np.uintp),
        #
        # End-index of the slice containing the node's training samples indexes
        ("end_train", np.uintp),
        #
        # Index of the first validation (out-of-the-bag) sample in the node. We have
        #   that partition_valid[start_valid:end_valid] contains the indexes of the
        #   node's validation samples
        ("start_valid", np.uintp),
        #
        # End-index of the slice containing the node's validation samples indexes
        ("end_valid", np.uintp),
    ]
)

# Convert this dtype to a numba-usable data-type
node_type = from_dtype(node_dtype)

# TODO: many attributes are not used in the node_context ?

# A node_context contains all the data required to find the best split of the node
node_context_type = [
    # This array contains the index of all the features. This will be modified
    # inplace when sampling features to be considered for splits
    ("features_pool", uintp[::1]),
    #
    # This array will contain the features sampled uniformly at random (without
    # replacement) to be considered for splits
    ("features_sampled", uintp[::1]),
    #
    # Do we need to perform features sampling ?
    ("sample_features", boolean),
    #
    # Number of training samples in the node
    ("n_samples_train", intp),
    #
    # Number of validation (out-of-the-bag) samples in the node
    ("n_samples_valid", intp),
    #
    # Weighted number of training samples in the node
    ("w_samples_train", float32),
    #
    # Weighted number of validation (out-of-the-bag) samples in the node
    ("w_samples_valid", float32),
    #
    # Index of the first training sample in the node. We have that
    #   partition_train[start_train:end_train] contains the indexes of the node's
    #   training samples
    ("start_train", intp),
    #
    # End-index of the slice containing the node's training samples indexes
    ("end_train", intp),
    #
    # Index of the first validation (out-of-the-bag) sample in the node. We have
    #   that partition_valid[start_valid:end_valid] contains the indexes of the
    #   node's validation samples
    ("start_valid", intp),
    #
    # End-index of the slice containing the node's validation samples indexes
    ("end_valid", intp),
    #
    # Validation loss of the node, computed on validation (out-of-the-bag) samples
    ("loss_valid", float32),
    #
    # Weighted number of training samples for each (feature, bin) in the node
    ("w_samples_train_in_bins", float32[:, ::1]),
    #
    # Weighted number of validation samples for each (feature, bin) in the node
    ("w_samples_valid_in_bins", float32[:, ::1]),
]

node_classifier_context_type = [
    *node_context_type,
    #
    # Weighted number of training samples for each (feature, bin, label) in the node
    ("y_sum", float32[:, :, ::1]),
    #
    # Prediction produced by the node using the training data it contains
    ("y_pred", float32[::1]),
]

node_regressor_context_type = [
    *node_context_type,
    #
    # Weighted sum of the labels for each (feature, bin) in the node
    ("y_sum", float32[:, ::1]),
    #
    # Weighted sum of the squared labels for each (feature, bin) in the node
    ("y_sq_sum", float32[:, ::1]),
    #
    # Prediction produced by the node using the training data it contains
    ("y_pred", float32),
]

@jitclass(node_classifier_context_type)
class NodeClassifierContext:
    """A node_context contains all the data required to find the best split of the
    node for a classification tree.

    Parameters
    ----------
    tree_context : TreeClassifierContext
        An object describing the context of the tree

    Attributes
    ----------
    features_pool : ndarray
        Array of shape (n_features,) of intp dtype containing all the features
        indexes, namely [0, ..., n_features-1]

    features_sampled : ndarray
        Array of shape (max_features,) of intp dtype containing sampled features to
        be considered for possible splitting

    sample_features : bool
        Do we need to perform features sampling ? (This if True if n_features !=
        max_features and False otherwise)

    features : ndarray
        Array of shape (max_features,) of intp dtype containing the candidate features
        for splitting. These features are chosen uniformly at random each time a
        split is looked for

    n_samples_train : int
        Number of training samples in the node

    n_samples_valid : int
        Number of validation (out-of-the-bag) samples in the node

    w_samples_train : float
        Weighted number of training samples in the node

    w_samples_valid : float
        Weighted number of validation (out-of-the-bag) samples in the node

    start_train : int
        Index of the first training sample in the node. We have that
        partition_train[start_train:end_train] contains the indexes of the node's
        training samples

    end_train : int
        End-index of the slice containing the node's training samples indexes

    start_valid : int
        Index of the first validation (out-of-the-bag) sample in the node. We have
        that partition_valid[start_valid:end_valid] contains the indexes of the
        node's validation samples

    end_valid : int
        End-index of the slice containing the node's validation samples indexes

    loss_valid : float
        Validation loss of the node, computed on validation (out-of-the-bag) samples

    w_samples_train_in_bins : ndarray
        Weighted number of training samples for each (feature, bin) in the node

    w_samples_valid_in_bins : ndarray
        Weighted number of validation samples for each (feature, bin) in the node

    y_sum : ndarray
        Weighted number of training samples for each (feature, bin, label) in the node

    y_pred : ndarray
        Prediction produced by the node using the training data it contains
    """

    def __init__(self, tree_context):
        init_node_context(tree_context, self)
        max_features = tree_context.max_features
        max_bins = tree_context.max_bins
        n_classes = tree_context.n_classes
        self.y_sum = np.empty((max_features, max_bins, n_classes), dtype=np.float32)
        self.y_pred = np.empty(n_classes, dtype=np.float32)

@jitclass(node_regressor_context_type)
class NodeRegressorContext:
    """A node_context contains all the data required to find the best split of the
    node for a regression tree.

    Parameters
    ----------
    tree_context : TreeRegressorContext
        An object describing the context of the tree

    Attributes
    ----------
    features_pool : ndarray
        Array of shape (n_features,) of intp dtype containing all the features
        indexes, namely [0, ..., n_features-1]

    features_sampled : ndarray
        Array of shape (max_features,) of intp dtype containing sampled features to
        be considered for possible splitting

    sample_features : bool
        Do we need to perform features sampling ? (This if True if n_features !=
        max_features and False otherwise)

    features : ndarray
        Array of shape (max_features,) of intp dtype containing the candidate features
        for splitting. These features are chosen uniformly at random each time a
        split is looked for

    n_samples_train : int
        Number of training samples in the node

    n_samples_valid : int
        Number of validation (out-of-the-bag) samples in the node

    w_samples_train : float
        Weighted number of training samples in the node

    w_samples_valid : float
        Weighted number of validation (out-of-the-bag) samples in the node

    start_train : int
        Index of the first training sample in the node. We have that
        partition_train[start_train:end_train] contains the indexes of the node's
        training samples

    end_train : int
        End-index of the slice containing the node's training samples indexes

    start_valid : int
        Index of the first validation (out-of-the-bag) sample in the node. We have
        that partition_valid[start_valid:end_valid] contains the indexes of the
        node's validation samples

    end_valid : int
        End-index of the slice containing the node's validation samples indexes

    loss_valid : float
        Validation loss of the node, computed on validation (out-of-the-bag) samples

    w_samples_train_in_bins : ndarray
        Weighted number of training samples for each (feature, bin) in the node

    w_samples_valid_in_bins : ndarray
        Weighted number of validation samples for each (feature, bin) in the node

    y_sum : ndarray
        Weighted sum of the labels for each (feature, bin) in the node

    y_sq_sum : ndarray
        Weighted sum of the squared labels for each (feature, bin) in the node

    y_pred : float
        Prediction produced by the node using the training data it contains
    """

    def __init__(self, tree_context):
        init_node_context(tree_context, self)
        max_features = tree_context.max_features
        max_bins = tree_context.max_bins
        self.y_sum = np.empty((max_features, max_bins), dtype=np.float32)
        self.y_sq_sum = np.empty((max_features, max_bins), dtype=np.float32)
        self.y_pred = 0.0

NodeClassifierContextType = get_type(NodeClassifierContext)
NodeRegressorContextType = get_type(NodeRegressorContext)

@jit(
    [
        void(TreeClassifierContextType, NodeClassifierContextType),
        void(TreeRegressorContextType, NodeRegressorContextType),
    ],
    nopython=True,
    nogil=True,
)
def init_node_context(tree_context, node_context):
    """A common initializer for NodeContextClassifier and NodeContextRegressor.
    Since numba's jitclass support for inheritance is inexistant, we use this
    function both in the constructors of NodeContextClassifier and NodeContextRegressor.

    Parameters
    ----------
    tree_context : TreeContext
        An object describing the context of the tree

    node_context : NodeClassifierContext or NodeRegressorContext
        The node context that we be partly initialized by this function
    """
    max_features = tree_context.max_features
    max_bins = tree_context.max_bins
    n_features = tree_context.n_features
    # If max_features is the same as n_features, we don't need to sample features
    node_context.sample_features = max_features != n_features
    # This array contains the index of all the features. This will be modified
    # inplace when sampling features to be considered for splits
    node_context.features_pool = np.arange(0, n_features, dtype=np.uintp)
    # This array will contain the features sampled uniformly at random (without
    # replacement) to be considered for splits
    node_context.features_sampled = np.arange(0, max_features, dtype=np.uintp)
    node_context.w_samples_train_in_bins = np.empty(
        (max_features, max_bins), dtype=np.float32
    )
    node_context.w_samples_valid_in_bins = np.empty(
        (max_features, max_bins), dtype=np.float32
    )

@jit(
    void(
        TreeClassifierContextType, NodeClassifierContextType, uintp, uintp, uintp, uintp
    ),
    nopython=True,
    nogil=True,
    boundscheck=False,
    locals={
        "w_samples_train_in_bins": float32[:, ::1],
        "w_samples_valid_in_bins": float32[:, ::1],
        "y_sum": float32[:, :, ::1],
        "y_pred": float32[::1],
        "features": uintp[::1],
        "X": uint8[:, :],
        "y": float32[::1],
        "sample_weights": float32[::1],
        "partition_train": uintp[::1],
        "partition_valid": uintp[::1],
        "n_classes": uintp,
        "dirichlet": float32,
        "train_indices": uintp[::1],
        "valid_indices": uintp[::1],
        "w_samples_train": float32,
        "w_samples_valid": float32,
        "f": uintp,
        "loss_valid": float32,
        "feature": uintp,
        "sample": uintp,
        "bin": uint8,
        "label": uintp,
        "sample_weight": float32,
        "k": uintp,
    },
)
def compute_node_classifier_context(
    tree_context, node_context, start_train, end_train, start_valid, end_valid
):
    """Computes the node context from the data and from the tree context for
    classification. Computations are saved in the passed node_context.

    Parameters
    ----------
    tree_context : TreeContext
        The tree context

    node_context : NodeClassifierContext
        The node context that this function will compute

    start_train : int
        Index of the first training sample in the node. We have that
        partition_train[start_train:end_train] contains the indexes of the node's
        training samples

    end_train : int
        End-index of the slice containing the node's training samples indexes

    start_valid : int
        Index of the first validation (out-of-the-bag) sample in the node. We have
        that partition_valid[start_valid:end_valid] contains the indexes of the
        node's validation samples

    end_valid : int
        End-index of the slice containing the node's validation samples indexes
    """
    # Initialize the things from the node context
    w_samples_train_in_bins = node_context.w_samples_train_in_bins
    w_samples_valid_in_bins = node_context.w_samples_valid_in_bins
    y_sum = node_context.y_sum
    y_pred = node_context.y_pred

    # If necessary, sample the features
    if node_context.sample_features:
        sample_without_replacement(
            node_context.features_pool, node_context.features_sampled
        )

    features = node_context.features_sampled
    w_samples_train_in_bins.fill(0.0)
    w_samples_valid_in_bins.fill(0.0)
    y_sum.fill(0.0)
    y_pred.fill(0.0)

    # Get information from the tree context
    X = tree_context.X
    y = tree_context.y
    sample_weights = tree_context.sample_weights
    partition_train = tree_context.partition_train
    partition_valid = tree_context.partition_valid
    n_classes = tree_context.n_classes
    dirichlet = tree_context.dirichlet

    # The indices of the training samples contained in the node
    train_indices = partition_train[start_train:end_train]
    valid_indices = partition_valid[start_valid:end_valid]

    # Weighted number of training and validation samples
    w_samples_train = 0.0
    w_samples_valid = 0.0

    # A counter for the features
    f = 0
    # The validation loss
    loss_valid = 0.0

    # TODO: unrolling the for loop could be faster
    # For-loop on features first and then samples (X is F-major)

    for feature in features:
        # Compute statistics about training samples
        for sample in train_indices:
            bin = X[sample, feature]
            label = uintp(y[sample])
            sample_weight = sample_weights[sample]
            if f == 0:
                w_samples_train += sample_weight
                y_pred[label] += sample_weight
            # One more sample in this bin for the current feature
            w_samples_train_in_bins[f, bin] += sample_weight
            # One more sample in this bin for the current feature with this label
            y_sum[f, bin, label] += sample_weight

        # TODO: we should put this outside so that we can change the dirichlet parameter
        # without rebuilding the tree
        # The prediction is given by the formula
        #   y_k = (n_k + dirichlet) / (n_samples + dirichlet * n_classes)
        # where n_k is the number of samples with label class k
        if f == 0:
            for k in range(n_classes):
                y_pred[k] = (y_pred[k] + dirichlet) / (
                    w_samples_train + n_classes * dirichlet
                )

        # Compute sample counts about validation samples
        for sample in valid_indices:
            bin = X[sample, feature]
            sample_weight = sample_weights[sample]
            if f == 0:
                w_samples_valid += sample_weight
                label = uintp(y[sample])
                # TODO: aggregation loss is hard-coded here. Call a function instead
                #  when implementing other losses
                loss_valid += - w_samples_valid * log(y_pred[label])

            w_samples_valid_in_bins[f, bin] += sample_weight

        f += 1

    # Save remaining things in the node context
    node_context.n_samples_train = end_train - start_train
    node_context.n_samples_valid = end_valid - start_valid
    node_context.loss_valid = loss_valid
    node_context.w_samples_train = w_samples_train
    node_context.w_samples_valid = w_samples_valid

@jit(
    void(
        TreeRegressorContextType, NodeRegressorContextType, uintp, uintp, uintp, uintp
    ),
    nopython=True,
    nogil=True,
    boundscheck=False,
    locals={
        "w_samples_train_in_bins": float32[:, ::1],
        "w_samples_valid_in_bins": float32[:, ::1],
        "y_sum": float32[:, ::1],
        "y_sq_sum": float32[:, ::1],
        "y_pred": float32,
        "features": uintp[::1],
        "X": uint8[:, :],
        "y": float32[::1],
        "sample_weights": float32[::1],
        "partition_train": uintp[::1],
        "partition_valid": uintp[::1],
        "train_indices": uintp[::1],
        "valid_indices": uintp[::1],
        "w_samples_train": float32,
        "w_samples_valid": float32,
        "f": uintp,
        "loss_valid": float32,
        "feature": uintp,
        "sample": uintp,
        "bin": uint8,
        "label": float32,
        "sample_weight": float32,
        "k": uintp,
        "w_y": float32,
    },
)
def compute_node_regressor_context(
    tree_context, node_context, start_train, end_train, start_valid, end_valid
):
    """Computes the node context from the data and from the tree context for regression.
    Computations are saved in the passed node_context.

    Parameters
    ----------
    tree_context : TreeContext
        The tree context

    node_context : NodeRegressorContext
        The node context that this function will compute

    start_train : int
        Index of the first training sample in the node. We have that
        partition_train[start_train:end_train] contains the indexes of the node's
        training samples

    end_train : int
        End-index of the slice containing the node's training samples indexes

    start_valid : int
        Index of the first validation (out-of-the-bag) sample in the node. We have
        that partition_valid[start_valid:end_valid] contains the indexes of the
        node's validation samples

    end_valid : int
        End-index of the slice containing the node's validation samples indexes
    """
    # Initialize the things from the node context
    w_samples_train_in_bins = node_context.w_samples_train_in_bins
    w_samples_valid_in_bins = node_context.w_samples_valid_in_bins
    y_sum = node_context.y_sum
    y_sq_sum = node_context.y_sq_sum

    # If necessary, sample the features
    if node_context.sample_features:
        sample_without_replacement(
            node_context.features_pool, node_context.features_sampled
        )

    features = node_context.features_sampled
    w_samples_train_in_bins.fill(0.0)
    w_samples_valid_in_bins.fill(0.0)
    y_sum.fill(0.0)
    y_sq_sum.fill(0.0)
    y_pred = 0.0

    # Get information from the tree context
    X = tree_context.X
    y = tree_context.y
    sample_weights = tree_context.sample_weights
    partition_train = tree_context.partition_train
    partition_valid = tree_context.partition_valid

    # The indices of the training samples contained in the node
    train_indices = partition_train[start_train:end_train]
    valid_indices = partition_valid[start_valid:end_valid]

    # Weighted number of training and validation samples
    w_samples_train = 0.0
    w_samples_valid = 0.0

    # A counter for the features
    f = 0
    # The validation loss
    loss_valid = 0.0

    # TODO: unrolling the for loop could be faster
    # For-loop on features first and then samples (X is F-major)

    for feature in features:
        # Compute statistics about training samples
        for sample in train_indices:
            bin = X[sample, feature]
            label = y[sample]
            sample_weight = sample_weights[sample]
            w_y = sample_weight * label
            if f == 0:
                w_samples_train += sample_weight
                y_pred += w_y
            # One more sample in this bin for the current feature
            w_samples_train_in_bins[f, bin] += sample_weight
            # One more sample in this bin for the current feature with this label

            y_sum[f, bin] += w_y
            y_sq_sum[f, bin] += w_y * label

        # The prediction is simply the weighted average of labels
        if f == 0:
            y_pred /= w_samples_train

        # Compute sample counts about validation samples
        for sample in valid_indices:
            bin = X[sample, feature]
            sample_weight = sample_weights[sample]
            if f == 0:
                w_samples_valid += sample_weight
                label = y[sample]
                # TODO: aggregation loss is hard-coded here. Call a function instead
                #  when implementing other losses
                loss_valid += sample_weight * (label - y_pred) * (label - y_pred)

            w_samples_valid_in_bins[f, bin] += sample_weight

        f += 1

    # Save remaining things in the node context
    node_context.y_pred = y_pred
    node_context.n_samples_train = end_train - start_train
    node_context.n_samples_valid = end_valid - start_valid
    node_context.loss_valid = loss_valid
    node_context.w_samples_train = w_samples_train
    node_context.w_samples_valid = w_samples_valid


"""Spectral shaping."""

from scipy.interpolate import interp1d

import numpy as np

# most of the contents of this file is heritage or not used.  That reflects
# a change in CGI practice.  Formerly, we used CBE values for coatings, etc.
# that posture has changed: we now lump things into a scalar value, which
# removes the spectral shaping from anything other than the star, detector,
# or FPM.

CGI_PUPIL_AREA = 0.641*2.363*2.363  # 64.1% of the enclosing rectangle (close enough)

# standard LOWFS bandpass, 9 wavelengths of spectral sampling
LOWFS_LOWER = (575-64)/1e3
LOWFS_UPPER = (575+64)/1e3

def mk_bandpass(n):
    """Make a LOWFS bandpass with n wavelengths."""
    if n == 1:
        return np.array([.575])

    return np.linspace(LOWFS_LOWER, LOWFS_UPPER, n)

LOWFS_BANDPASS = mk_bandpass(9)

CGI_BASIC_COATINGS = {
    'HRC': 7,       # = OTA+TCA, per Bijan
    'FSS99': 7+1,   # = CGI front-end + rad shield mirror (LOCAM)
    'Al': 2,        # = DM1+DM2, CGI front-end
    'BBAR': 3,      # LOWFS lenses
}
# additions...
# HLC => +1 FSS99 (SPM)
# SPC => +1 Al (SPM)

class ThroughputDatabase:
    """Database of spectral throughputs."""

    def __init__(self, emccd):  # historical: first arg coatings
        """Create a new database of throughputs.

        Parameters
        ----------
        coatings : `dict`
            keys of strings, which are the names of coatings
            values of scipy.interpolation.interp1d returns
        emccd : `scipy.interpolation.interp1d`
            interp1d of emccd data

        """
        # self.coatings = coatings
        self.emccd_interpf = emccd

    def __call__(self, corongraph_mode, wavelengths):
        """Spectral shape andthroughput for mode & wavelengths."""
        # hrc = self.coatings['HRC']
        # fss99 = self.coatings['FSS99']
        # al = self.coatings['Al']
        # bbar = self.coatings['bbar']

        # n_hrc = CGI_BASIC_COATINGS['HRC']
        # n_fss99 = CGI_BASIC_COATINGS['FSS99']
        # n_al = CGI_BASIC_COATINGS['Al']
        # n_bbar = CGI_BASIC_COATINGS['BBAR']
        # if corongraph_mode == 'HLC':
        #     #  not += 1 to avoid any chance of modifying the dict
        #     n_fss99 = n_fss99 + 1
        # else:
        #     n_al = n_al + 1

        # hrc = hrc(wavelengths) ** n_hrc
        # fss99 = fss99(wavelengths) ** n_fss99
        # al = al(wavelengths) ** n_al
        # bbar = bbar(wavelengths) ** n_bbar
        emccd = self.emccd_interpf(wavelengths)
        # return hrc * fss99 * al * bbar * emccd
        corongraph_mode = corongraph_mode.upper()
        # from LOWFS_photometry 20200801, Throughput sheet.  BOL REQs
        if corongraph_mode == 'HLC':
            return 0.3443 * emccd
        elif corongraph_mode == 'SPEC':
            return 0.3092 * emccd
        elif corongraph_mode == 'WFOV':
            return 0.3092 * emccd
        else:
            raise ValueError('invalid coronagraph mode: must be HLC, SPEC, or WFOV')

    @classmethod
    def bijan_data(cls, data_fldr):
        """Load data from Bijan."""
        emccd_data = np.genfromtxt(data_fldr/'emccd-spectral-qe.csv', delimiter=',', skip_header=1)
        emccd_wvl = emccd_data[:, 0]/1e3
        emccd_dat = emccd_data[:, 1]
        emccd_interpf = interp1d(emccd_wvl, emccd_dat, bounds_error=False, fill_value=0)

        # coating_data = np.genfromtxt(data_fldr/'coatings.csv', delimiter=',', skip_header=1)
        # coating_wvl = coating_data[:, 0]
        # coatings = {}
        # for i, lbl in enumerate(['HRC', 'FSS99', 'Al', 'BBAR']):
        #     data = coating_data[:, i+1]
        #     interpf = interp1d(coating_wvl, data)
        #     coatings[lbl] = interpf

        # return ThroughputDatabase(coatings, emccd_interpf)
        return ThroughputDatabase(emccd_interpf)

class StellarDatabase:
    """Database of stellar spectra."""

    def __init__(self, stellar_table, stellar_labels):
        """Create a new database of stellar throughputs.

        Parameters
        ----------
        stellar_table : `numpy.ndarray`
            columnnar table of stellar spectral fluxes
        stellar_labels : `numpy.ndarray`
            labels associated with the columns, as [g0v, av5], and so forth.
            The index of the label identifies its column

        """
        self.stellar_interpfs = {}
        self.table = stellar_table
        self.stellar_labels = stellar_labels

    def __call__(self, star_type, wavelengths, lower_lim=-1, upper_lim=-1):
        """Spectral weights associated with a given star type at a given wavelength."""
        stellar_interpf = self.stellar_interpfs.get(star_type, None)
        if lower_lim == -1:
            lower_lim = LOWFS_LOWER
        if upper_lim == -1:
            upper_lim = LOWFS_UPPER

        if stellar_interpf is None:
            # this line throws value error if star_type not in list
            i = self.stellar_labels.index(star_type)

            # table has units of W/(m^2 m) == J/(s m^2 m)
            # => divide by J/ph, multiply by density, multiply by area
            # == ph/s.
            # this could be cleaner with astropy units, but c'est la vie
            wvl = self.table[:, 0]*1e6  # 1e6; m => um
            photon_energy = self.table[:, 1]
            dat = self.table[:, i]
            dat = dat / photon_energy * 1e-9 * CGI_PUPIL_AREA  # 1e-9 = sampling density, 1 nm
            stellar_interpf = interp1d(wvl, dat)
            self.stellar_interpfs[star_type] = stellar_interpf

        ret = stellar_interpf(wavelengths)
        ret[wavelengths < lower_lim] = 0
        ret[wavelengths > upper_lim] = 0
        return ret

    def sparsity_fudge_factor(self, star_type, wavelengths, lower_lim=-1, upper_lim=-1):
        """Fudge factor to deal with spectral sparsity for a given set of wavelengths."""
        true = self(star_type, wavelengths, lower_lim=lower_lim, upper_lim=upper_lim).sum()
        wvl_dense = self.table[:, 0] * 1e6
        dense = self(star_type, wvl_dense, lower_lim=lower_lim, upper_lim=upper_lim).sum()
        return dense/true

    @classmethod
    def bijan_data(cls, data_fldr):
        """Load data from Bijan."""
        stellar_data = np.genfromtxt(data_fldr/'stellar-spectra.csv', delimiter=',', skip_header=1)
        labels = [
            'wvl',
            'photon energy',
            'b3v',
            'a0v',
            'a5v',
            'f5v',
            'g0v',
            'f5v',
            'k0v',
            'k5v',
            'm0v',
            'm5v'
        ]

        return StellarDatabase(stellar_data, labels)


# This source code is part of the Hydride package and is distributed
# under the 3-Clause BSD License. Please see 'LICENSE.rst' for further
# information.

__name__ = "hydride"
__author__ = "Patrick Kunzmann"
__all__ = ["FragmentLibrary"]

from os.path import join, dirname, abspath
import warnings
import pickle
from biotite.structure.error import BadStructureError
from biotite.structure import BondType, displacement
import numpy as np

class FragmentLibrary:
    """
    A molecule fragment library for estimation of hydrogen positions.

    For each molecule added to the :class:`FragmentLibrary`,
    the molecule is split into fragments.
    Each fragment consists of

        - A central heavy atom,
        - bond order and position of its bonded heavy atoms and
        - and positions of bonded hydrogen atoms.

    The properties of the fragment (central atom element,
    central atom charge, order of connected bonds) are stored in
    a dictionary mapping these properties to heavy and hydrogen atom
    positions.

    If hydrogen atoms should be added to a target structure,
    the target structure is also split into fragments.
    Now the corresponding reference fragment in the library dictionary
    is accessed for each fragment.
    The corresponding atom coordinates of the reference fragment
    are superimposed [1]_ [2]_ onto the target fragment to obtain the
    hydrogen coordinates for the heavy atom.

    The constructor of this class creates an empty library.

    References
    ----------

    .. [1] W Kabsch,
       "A solution for the best rotation to relate two sets of vectors."
       Acta Cryst, 32, 922-923 (1976).

    .. [2] W Kabsch,
       "A discussion of the solution for the best rotation to relate
       two sets of vectors."
       Acta Cryst, 34, 827-828 (1978).
    """

    _std_library = None

    def __init__(self):
        self._frag_dict = {}

    @staticmethod
    def standard_library():
        """
        Get the standard :class:`FragmentLibrary`.
        The library contains fragments from all molecules in the
        *RCSB* *Chemical Component Dictionary*.

        Returns
        -------
        library : FragmentLibrary
            The standard library.
        """
        if FragmentLibrary._std_library is None:
            FragmentLibrary._std_library = FragmentLibrary()
            file_name = join(dirname(abspath(__file__)), "fragments.pickle")
            with open(file_name, "rb") as fragments_file:
                FragmentLibrary._std_library._frag_dict \
                    = pickle.load(fragments_file)
        return FragmentLibrary._std_library

    def add_molecule(self, molecule):
        """
        Add the fragments of a molecule to the library.

        Parameters
        ----------
        molecule : AtomArray
            A molecule, whose fragments should be added to the library.
            The structure must contain hydrogen atoms for each
            applicable heavy atom.
            The molecule must have an associated :class:`BondList`.
            The molecule must also include the *charge* annotation
            array, depicting the formal charge for each atom.
        """
        fragments = _fragment(molecule)
        for i, fragment in enumerate(fragments):
            if fragment is None:
                continue
            (
                central_element, central_charge, stereo, bond_types,
                center_coord, heavy_coord, hydrogen_coord
            ) = fragment
            # Translate the coordinates,
            # so the central heavy atom is at origin
            centered_heavy_coord = heavy_coord - center_coord
            centered_hydrogen_coord = hydrogen_coord - center_coord
            self._frag_dict[
                (central_element, central_charge, stereo, tuple(bond_types))
            ] = (
                # Information about the origin of the fragment
                # # for debugging purposes
                molecule.res_name[i], molecule.atom_name[i],
                # The interesting information
                centered_heavy_coord, centered_hydrogen_coord
            )
            if stereo != 0:
                # Also include the opposite enantiomer in the library
                # by reflecting the coordinates along an arbitrary axis
                stereo *= -1
                refl_centered_heavy_coord = centered_heavy_coord.copy()
                refl_centered_hydrogen_coord = centered_hydrogen_coord.copy()
                refl_centered_heavy_coord[..., 0] *= -1
                refl_centered_hydrogen_coord[..., 0] *= -1
                self._frag_dict[(
                    central_element, central_charge, stereo, tuple(bond_types)
                )] = (
                    molecule.res_name[i], molecule.atom_name[i],
                    refl_centered_heavy_coord, refl_centered_hydrogen_coord
                )

    def calculate_hydrogen_coord(self, atoms, mask=None, box=None):
        """
        Estimate the hydrogen coordinates for each atom in a given
        structure/molecule.

        Parameters
        ----------
        atoms : AtomArray, shape=(n,)
            The structure to get the hydrogen atom positions for.
            The structure must not contain any hydrogen atoms.
            The structure must have an associated :class:`BondList`.
            The structure must also include the *charge* annotation
            array, depicting the formal charge for each atom.
        mask : ndarray, shape=(n,), dtype=bool
            A boolean mask that is true for each heavy atom, where
            corresponsing hydrogen atom positions should be calculated.
            By default, hydrogen atoms are calculated for all applicable
            atoms.
        box : bool or array-like, shape=(3,3), dtype=float, optional
            If this parameter is set, periodic boundary conditions are
            taken into account (minimum-image convention), based on
            the box vectors given with this parameter.
            If `box` is set to true, the box vectors are taken from the
            ``box`` attribute of `atoms` instead.

        Returns
        -------
        hydrogen_coord : list of (ndarray, shape=(k,3), dtype=np.float32), length=n
            A list of hydrogen coordinates for each atom in the input
            `atoms`.
            *k* is the number of hydrogen atoms for this atom.
            Each atom, that is not included in the input `mask` or which
            has no fitting fragment in the library,
            has no corresponding hydrogen coordinates, i.e. *k* is 0.
        """
        if mask is None:
            mask = np.ones(atoms.array_length(), dtype=bool)

        # The target and reference heavy atom coordinates
        # for each fragment
        tar_frag_center_coord = np.zeros(
            (atoms.array_length(), 3), dtype=np.float32
        )
        tar_frag_heavy_coord = np.zeros(
            (atoms.array_length(), 3, 3), dtype=np.float32
        )
        ref_frag_heavy_coord = np.zeros(
            (atoms.array_length(), 3, 3), dtype=np.float32
        )
        # The amount of hydrogens varies for each fragment
        # -> padding with NaN
        # The maximum number of bond hydrogen atoms is 4
        ref_frag_hydrogen_coord = np.full(
            (atoms.array_length(), 4, 3), np.nan, dtype=np.float32
        )

        # Fill the coordinate arrays
        fragments = _fragment(atoms, mask)
        for i, fragment in enumerate(fragments):
            if fragment is None:
                # This atom is not in mask
                continue
            (
                central_element, central_charge, stereo, bond_types,
                center_coord, heavy_coord, _
            ) = fragment
            tar_frag_center_coord[i] = center_coord
            tar_frag_heavy_coord[i] = heavy_coord
            # The hydrogen_coord can be ignored:
            # In the target structure are no hydrogen atoms
            hit = self._frag_dict.get(
                (central_element, central_charge, stereo, tuple(bond_types))
            )
            if hit is None:
                warnings.warn(
                    f"Missing fragment for atom '{atoms.atom_name[i]}' "
                    f"at position {i}"
                )
            else:
                _, _, ref_heavy_coord, ref_hydrogen_coord = hit
                ref_frag_heavy_coord[i] = ref_heavy_coord
                ref_frag_hydrogen_coord[i, :len(ref_hydrogen_coord)] \
                    = ref_hydrogen_coord

        # Translate the target coordinates,
        # so the central heavy atom is at origin
        # This has already been done for the reference atoms
        # in the 'add_molecule()' method
        if box is not None:
            # Find shortest possible displacement vector for each heavy
            # atom according to minimum image convention
            if box is True:
                if atoms.box is None:
                    raise ValueError("Input structure has no associated box")
                box = atoms.box
            else:
                # Box vectors are given as array-like object
                box = np.asarray(box)
        tar_frag_heavy_coord = displacement(
            tar_frag_center_coord[:, np.newaxis, :], tar_frag_heavy_coord,
            box
        )

        # Get the rotation matrix required for superimposition of
        # the reference coord to the target coord
        matrices = _get_rotation_matrices(
            tar_frag_heavy_coord, ref_frag_heavy_coord
        )
        # Rotate the reference hydrogen atoms, so they fit the
        # target heavy atoms
        tar_frag_hydrogen_coord = _rotate(ref_frag_hydrogen_coord, matrices)
        # Translate hydrogen atoms to the position of the
        # non-centered central heavy target atom
        tar_frag_hydrogen_coord += tar_frag_center_coord[:, np.newaxis, :]

        # Turn into list and remove NaN paddings
        tar_frag_hydrogen_coord = [
            # If the x-coordinate is NaN it is expected that
            # y and z are also NaN
            coord[~np.isnan(coord[:, 0])] for coord in tar_frag_hydrogen_coord
        ]

        return tar_frag_hydrogen_coord

def _fragment(atoms, mask=None):
    """
    Create fragments for the input structure/molecule.

    Parameters
    ----------
    atoms : AtomArray, shape=(n,)
        The structure to be fragmented.
        The structure must have an associated :class:`BondList`.
        The structure must also include the *charge* annotation
        array, depicting the formal charge for each atom.
    mask : ndarray, shape=(n,), dtype=bool
        A boolean mask that is true for each heavy atom for which a
        fragment should be created.

    Returns
    -------
    fragments : list of tuple(str, int, int, ndarray, ndarray, ndarray), length=n
        The fragments.
        The tuple elements are

            #. the central atom element,
            #. the central atom charge,
            #. the enantiomer for stereocenter
               (``-1`` and ``1`` based on a custom nomenclature),
            #. the :class:`BondType` for each bonded heavy atom,
            #. the coordinates of the central atom,
            #. 3 coordinates of bonded heavy atoms (includes padding
               values, if there are not enough heavy atoms),
            #. the coordinates of bonded hydrogen atoms.

        ``None`` for each atom not included by the `mask`.
    """
    if mask is None:
        mask = np.ones(atoms.array_length(), dtype=bool)

    if atoms.bonds is None:
        raise BadStructureError(
            "The input structure must have an associated BondList"
        )

    fragments = [None] * atoms.array_length()

    all_bond_indices, all_bond_types = atoms.bonds.get_all_bonds()
    elements = atoms.element
    charges = atoms.charge
    coord = atoms.coord

    for i in range(atoms.array_length()):
        if not mask[i]:
            continue

        if elements[i] == "H":
            # Only create fragments for heavy atoms
            continue
        bond_indices = all_bond_indices[i]
        bond_types = all_bond_types[i]
        bond_indices = bond_indices[bond_indices != -1]
        bond_types = bond_types[bond_types != -1]

        heavy_mask = (elements[bond_indices] != "H")
        heavy_indices = bond_indices[heavy_mask]
        heavy_types = bond_types[heavy_mask]
        if (heavy_types == BondType.ANY).any():
            warnings.warn(
                f"Atom '{atoms.atom_name[i]}' in '{atoms.res_name[i]}' has an "
                f"undefined bond type and is ignored"
            )
            continue

        # Order the bonded atoms by their bond types
        # to remove atom order dependency in the matching step
        order = np.argsort(heavy_types)
        heavy_indices = heavy_indices[order]
        heavy_types = heavy_types[order]

        hydrogen_mask = ~heavy_mask
        hydrogen_coord = coord[bond_indices[hydrogen_mask]]

        # Special handling of nitrogen as central atom:
        # There are cases where the free electron pair can form
        # a partial double bond.
        # Although the bond order is formally 1 in this case,
        # it would enforce planar hydrogen positionioning
        # Therefore, a partial double bond is handled as bond type 7
        if elements[i] == "N":
            for j, remote_index in enumerate(heavy_indices):
                if heavy_types[j] != 1:
                    # This handling only applies to single bonds
                    continue
                rem_bond_indices = all_bond_indices[remote_index]
                rem_bond_indices = rem_bond_indices[rem_bond_indices != -1]
                rem_bond_types = all_bond_types[remote_index]
                rem_bond_types = rem_bond_types[rem_bond_types != -1]
                for rem_rem_index, bond_type in zip(
                    rem_bond_indices, rem_bond_types
                ):
                    # If the adjacent atom has a double bond
                    # the partial double bond condition is fulfilled
                    if bond_type == BondType.AROMATIC_DOUBLE or \
                       bond_type == BondType.DOUBLE:
                            heavy_types[j] = 7

        n_heavy_bonds = np.count_nonzero(heavy_mask)
        if n_heavy_bonds == 0:
            # The orientation is arbitrary
            # -> The fragment coord is the coord of the central atom
            # 4 times repeated
            heavy_coord = np.repeat(coord[np.newaxis, i, :], 3, axis=0)
            stereo = 0
        elif n_heavy_bonds == 1:
            # Include one atom further away
            # to get an unambiguous fragment
            remote_index = heavy_indices[0]
            rem_bond_indices = all_bond_indices[remote_index]
            rem_bond_indices = rem_bond_indices[rem_bond_indices != -1]
            rem_heavy_mask = (elements[rem_bond_indices] != "H")
            rem_heavy_indices = rem_bond_indices[rem_heavy_mask]
            # Use the coord of any heavy atom bonded to the remote
            # atom
            rem_rem_index = rem_heavy_indices[0]
            # Include the directly bonded atom two times, to give it a
            # greater weight in superimposition
            heavy_coord = coord[[remote_index, remote_index, rem_rem_index]]
            stereo = 0
        elif n_heavy_bonds == 2:
            heavy_coord = coord[[heavy_indices[0], heavy_indices[1], i]]
            stereo = 0
        elif n_heavy_bonds == 3:
            heavy_coord = coord[heavy_indices]
            center = coord[i]
            # Determine the enantiomer of this stereocenter
            # For performance reasons, the result does not follow the
            # R/S nomenclature, but a custom -1/1 based one, which also
            # unambiguously identifies the enantiomer
            n = np.cross(heavy_coord[0] - center, heavy_coord[1] - center)
            stereo = int(np.sign(np.dot(heavy_coord[2] - center, n)))
        elif n_heavy_bonds == 4:
            # The fragment is irrelevant, as there is no bonded hydrogen
            # -> The fragment coord is the coord of the central atom
            # 4 times repeated
            heavy_coord = np.repeat(coord[np.newaxis, i, :], 3, axis=0)
            stereo = 0
        else:
            warnings.warn(
                f"Atom '{atoms.atom_name[i]}' in "
                f"'{atoms.res_name[i]}' has more than 4 bonds to "
                f"heavy atoms ({n_heavy_bonds}) and is ignored"
            )
            heavy_coord = np.repeat(coord[np.newaxis, i, :], 3, axis=0)
            hydrogen_coord = np.zeros((0, 3), dtype=np.float32)
            stereo = 0
        central_coord = coord[i]
        fragments[i] = (
            elements[i], charges[i], stereo, heavy_types,
            central_coord, heavy_coord, hydrogen_coord
        )
    return fragments

def _get_rotation_matrices(fixed, mobile):
    """
    Get the rotation matrices to superimpose the given mobile
    coordinates into the given fixed coordinates, minimizing the RMSD.

    Uses the *Kabsch* algorithm.

    Parameters
    ----------
    fixed : ndarray, shape=(m,n,3), dtype=np.float32
        The fixed coordinates.
    mobile : ndarray, shape=(m,n,3), dtype=np.float32
        The mobile coordinates.

    Returns
    -------
    matrices : ndarray, shape=(m,3,3), dtype=np.float32
        The rotation matrices.
    """
    # Calculate cross-covariance matrices
    cov = np.sum(fixed[:,:,:,np.newaxis] * mobile[:,:,np.newaxis,:], axis=1)
    v, s, w = np.linalg.svd(cov)
    # Remove possibility of reflected atom coordinates
    reflected_mask = (np.linalg.det(v) * np.linalg.det(w) < 0)
    v[reflected_mask, :, -1] *= -1
    matrices = np.matmul(v, w)
    return matrices

def _rotate(coord, matrices):
    """
    Apply a rotation on given coordinates.

    Parameters
    ----------
    coord : ndarray, shape=(m,n,3), dtype=np.float32
        The coordinates.
    matrices : ndarray, shape=(m,3,3), dtype=np.float32
        The rotation matrices.
    """
    return np.transpose(
        np.matmul(matrices, np.transpose(coord, axes=(0, 2, 1))),
        axes=(0, 2, 1)
    )