#! /usr/bin/env python3

import random
import copy
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks_cwt

LENGTH = 50000000
N_TRIALS = 350
selection_strength = 1

replicates = 3000
plot_all = False
window = 1000000
step = 10000

out_pre = '/Volumes/Jacob_2TB_storage/sim_sec_recombination_mapping/genomics_scripts/analysis/'

class Sample_pool():
    """Holds samples, which need to be Sequence objects"""

    def __init__(self):
        self.pool = {}
        self.keys = []

    def add_sample(self, sequence):
        if self.keys:
            newkey = max(self.keys) + 1
        else:
            newkey = 0
        self.pool[newkey] = sequence
        self.keys = list(self.pool.keys())

    def pick_sample(self):
        """Selects a random sample"""
        sample = random.choice(self.keys)
        return self.pool[sample]

    def sample_freq(self, position):
        """Gathers a sample allele frequency for a given position"""
        reps = random.randint(36,40)
        mel = 0
        sim = 0
        sec = 0
        for i in range(reps):
            sample = self.pick_sample()
            if sample[position] == 1:
                mel += 1
            elif sample[position] == 2:
                sim += 1
            elif sample[position] == 3:
                sec += 1

        melf = mel / (mel + sim + sec)
        simf = sim / (mel + sim + sec)
        secf = sec / (mel + sim + sec)

        return melf, simf, secf

def bool_chance(num):
    x = random.random()
    if x < num:
        return True
    else:
        return False

def generate_parents():
    """
    Generate a set of 3 parents where all positions are 0, if mel specific 1,
    if sim specific 2, and if sec specific 3.
    Return Dicitonary of parents
    """

    everyN = 6000

    # get a list of the positions that would be differentiating spots
    id_spots = [random.randint(0,LENGTH) for i in range(int(LENGTH/everyN))]
    id_spots.sort()

    # make a list of the same length for each species
    mel_seq = [1] * len(id_spots)
    sim_seq = [2] * len(id_spots)
    sec_seq = [3] * len(id_spots)

    parents = {
        'mel' : mel_seq,
        'sim' : sim_seq,
        'sec' : sec_seq,

    }

    return parents, id_spots

def pick_selected_index(id_spots):
    """
    Takes the list of different sites between the three sequences
    Picks a random spot in the middle half of the sequences to be the selected site
    Returns the index of the differentiated spots list that is closest to the
        selected point, and the selected point
    """

    q1 = int(LENGTH / 8)
    q3 = int(7 * LENGTH / 8)
    sel_spot = random.randint(q1, q3)

    hold = 0
    for idx, x in enumerate(id_spots):
        if x < sel_spot:
            hold = x
        else:
            if (sel_spot - hold) <= (x - sel_spot):
                return (idx-1), sel_spot
            else:
                return idx, sel_spot

def reorder(one, two):
    if one >= two:
        return two, one
    else:
        return one, two

def gamma_model():
    """Compute the gamma distrubtion model for recombination.
    Return a dictionary with 100 steps (should never need that many)
    """
    model = {}
    prev = 0
    for i in range(1,101):
        # the computation:
            # (5*1morgan) (sum_to_i)(
            #   1 / (5i-1) * 1M^(5i) * dist^(5i-1) * e^(-(1M)*dist)
            # )
            # where distance is in morgans
        prev += ((1**(5*i)) * ((0.01*i)**((5*i)-1)) * (np.e ** (-(0.01*i)))) / (5*(i) - 1)
        model[i] = 5 * prev
    return model

def recombine(seq1, seq2, model):
    """Make a recombinant sequence from sequence 1 and 2"""

    idxs = [idx for idx, x in enumerate(seq1)]

    co1 = random.choice(idxs)

    half = int(len(idxs) / 2)

    if co1 < half:
        co2 = random.choice(idxs[half:])
    else:
        co2 = random.choice(idxs[:half])

    co1, co2 = reorder(co1, co2)

    if bool_chance(0.5):
        rec_seq = seq1[:co1] + seq2[co1:co2] + seq1[co2:]
    else:
        rec_seq = seq2[:co1] + seq1[co1:co2] + seq2[co2:]

    return rec_seq

def filter(seq, sel_idx, sel_strength):
    """Checks if a sequence contains the selected site"""
    if seq[sel_idx] == 2:
        return True
    else:
        if bool_chance(1-sel_strength):
            return True
        else:
            return False

def make_pool(parents, sel_idx, sex, sel_strength, model):
    """
    Input: mel, sim, sec parent sequence  , site to be selected on
    Returns: sample pool of recombinant sequences and mel sequences
    """

    pool = Sample_pool()
    for i in range(N_TRIALS):
        survive = False
        while not survive:
            recseq = recombine(parents['sim'], parents['sec'], model)
            survive = filter(recseq, sel_idx, sel_strength)
            if sex == 'female':
                survive = True
        pool.add_sample(recseq)
        # because diploids, we add one mel sequence for each recombinant sequence
        pool.add_sample(parents['mel'])

    return pool

def simulate_sequencing(pool, id_spots):
    table = []
    for idx, x in enumerate(id_spots):
        mel, sim, sec = pool.sample_freq(idx)
        table.append([x,mel,sim,sec])
    return table

def sex_difference(male_reps, female_reps):
    reps = {}

    for i in male_reps:
        reps[i] = []
        for idx, entry in enumerate(male_reps[i]):
            male = male_reps[i][idx]
            female = female_reps[i][idx]
            freq = [(male[0]), (male[1] - female[1]), (male[2] - female[2]), (male[3] - female[3])]
            reps[i].append(freq)

    return reps

def average_replicates(reps):
    table = []
    for pos, lis in enumerate(reps[0]):
        melav = (reps[0][pos][1] + reps[1][pos][1] + reps[2][pos][1]) / 3
        simav = (reps[0][pos][2] + reps[1][pos][2] + reps[2][pos][2]) / 3
        secav = (reps[0][pos][3] + reps[1][pos][3] + reps[2][pos][3]) / 3
        table.append([reps[0][pos][0], melav, simav, secav])
    return table

def rolling_average(table, window, step):

    win2 = window/2
    pos = window/2

    window_table = []
    posits = [x[0] for x in table]
    while pos < max(posits):
        melav = []
        simav = []
        secav = []
        for x in table:
            if (x[0] > pos-win2):
                if (x[0] < pos+win2):
                    melav.append(x[1])
                    simav.append(x[2])
                    secav.append(x[3])
                else:
                    break
        if melav and simav and secav:
            window_table.append([pos, np.mean(melav), np.mean(simav), np.mean(secav)])
        pos += step

    newtable = [['position','mel','sim','sec']]
    for x in window_table:
        newtable.append(x)

    return newtable

def estimate_max(table):
    table = table[1:]
    sim_freqs = [(x[0],(x[2]-x[3]))for x in table]
    sim_freqs = sorted(sim_freqs, key=lambda x: x[1])

    # find the peaks
    xs = [x[1] for x in sim_freqs]
    peaks = list(find_peaks_cwt(xs, np.arange(50, 200)))

    # this produces a list. Find the biggest one in the list
    big = (0,0)
    for peak in peaks:
        if sim_freqs[peak][1] > big[1]:
            big = (sim_freqs[peak][0], sim_freqs[peak][1])

    return big[0]

def write_to_csv(table, filename):
    with open(filename, 'w') as f:
        for line in table:
            line = [str(x) for x in line]
            f.write(','.join(line))
            f.write('\n')

def plot_frequencies(filename, sel_spot, esitmate, num):
    df = pd.read_csv(filename)
    fig = plt.figure(figsize=(6, 8))
    plt.plot(df['position'], df['mel'], color = 'blue', label = 'D.mel')
    plt.plot(df['position'], df['sim'], color = 'orange', label = 'D.sim')
    plt.plot(df['position'], df['sec'], color = 'red', label = 'D.sec')
    plt.axvline(x=sel_spot, color='black', label = 'actual site')
    plt.axvline(x=esitmate, color='green', label = 'estimated site')
    plt.ylim(-0.4,0.4)
    plt.legend()
    plt.ylabel('Allele Frequency (Male - Female)')
    plt.xlabel('Genomic position')
    plotname = filename.split('.csv')[0] + str(num) + '.pdf'
    plt.savefig(plotname)

parents, id_spots = generate_parents()
model = gamma_model()
differences = []
for i2 in range(replicates):
    sel_idx, sel_spot = pick_selected_index(id_spots)
    male_reps = {}
    female_reps = {}
    reps = {}
    # replicate 3 times
    for i in range(3):
        male_pool = make_pool(parents, sel_idx, 'male', selection_strength, model)
        female_pool = make_pool(parents, sel_idx, 'female', selection_strength, model)
        male_reps[i] = simulate_sequencing(male_pool, id_spots)
        female_reps[i] = simulate_sequencing(female_pool, id_spots)
    reps = sex_difference(male_reps, female_reps)
    table = average_replicates(reps)
    table = rolling_average(table, window, step)
    estimated_site = estimate_max(table)
    differences.append(estimated_site-sel_spot)
    if plot_all:
        outtable = out_pre + str(i2) + 'simulation_out.csv'
        write_to_csv(table, outtable)
        plot_frequencies(outtable, sel_spot, estimated_site, i2)

print(differences)
print("2x std:", np.std(differences))
print("Average:", np.average(differences))
plt.hist(differences, bins=20)
plt.savefig(out_pre+'confidence_hist.pdf')

outtable = out_pre + 'final_simulation_out.csv'
write_to_csv(table, outtable)
plot_frequencies(outtable, sel_spot, estimated_site, 'final')


#!/usr/bin/python
# -*- coding: utf-8 -*-

################################################################################
#
#   RMG - Reaction Mechanism Generator
#
#   Copyright (c) 2002-2010 Prof. William H. Green (whgreen@mit.edu) and the
#   RMG Team (rmg_dev@mit.edu)
#
#   Permission is hereby granted, free of charge, to any person obtaining a
#   copy of this software and associated documentation files (the 'Software'),
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense,
#   and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#   DEALINGS IN THE SOFTWARE.
#
################################################################################

"""

"""

import os
import os.path
import math
import logging
import numpy
from copy import copy, deepcopy

from base import Database, Entry, makeLogicNode

from rmgpy.quantity import constants
from rmgpy.thermo import *
from rmgpy.molecule import Molecule, Atom, Bond, Group

################################################################################

def saveEntry(f, entry):
    """
    Write a Pythonic string representation of the given `entry` in the thermo
    database to the file object `f`.
    """

    f.write('entry(\n')
    f.write('    index = {0:d},\n'.format(entry.index))
    f.write('    label = "{0}",\n'.format(entry.label))

    if isinstance(entry.item, Molecule):
        f.write('    molecule = \n')
        f.write('"""\n')
        f.write(entry.item.toAdjacencyList(removeH=True))
        f.write('""",\n')
    elif isinstance(entry.item, Group):
        f.write('    group = \n')
        f.write('"""\n')
        f.write(entry.item.toAdjacencyList())
        f.write('""",\n')
    else:
        f.write('    group = "{0}",\n'.format(entry.item))

    if isinstance(entry.data, ThermoData):
        f.write('    thermo = ThermoData(\n')
        f.write('        Tdata = {0!r},\n'.format(entry.data.Tdata))
        f.write('        Cpdata = {0!r},\n'.format(entry.data.Cpdata))
        f.write('        H298 = {0!r},\n'.format(entry.data.H298))
        f.write('        S298 = {0!r},\n'.format(entry.data.S298))
        if entry.data.Tmin is not None: f.write('        Tmin = {0!r},\n'.format(entry.data.Tmin))
        if entry.data.Tmax is not None: f.write('        Tmax = {0!r},\n'.format(entry.data.Tmax))
        f.write('    ),\n')
    elif isinstance(entry.data, Wilhoit):
        f.write('    thermo = Wilhoit(\n')
        f.write('        cp0 = {0!r},\n'.format(entry.data.cp0))
        f.write('        cpInf = {0!r},\n'.format(entry.data.cpInf))
        f.write('        a0 = {0:g},\n'.format(entry.data.a0))
        f.write('        a1 = {0:g},\n'.format(entry.data.a1))
        f.write('        a2 = {0:g},\n'.format(entry.data.a2))
        f.write('        a3 = {0:g},\n'.format(entry.data.a3))
        f.write('        B = {0!r},\n'.format(entry.data.B))
        f.write('        H0 = {0!r},\n'.format(entry.data.H0))
        f.write('        S0 = {0!r},\n'.format(entry.data.S0))
        if entry.data.Tmin is not None: f.write('        Tmin = {0!r},\n'.format(entry.data.Tmin))
        if entry.data.Tmax is not None: f.write('        Tmax = {0!r},\n'.format(entry.data.Tmax))
        f.write('    ),\n')
    elif isinstance(entry.data, MultiNASA):
        f.write('    thermo = MultiNASA(\n')
        f.write('        polynomials = [\n')
        for poly in entry.data.polynomials:
            f.write('            {0!r},\n'.format(poly))
        f.write('        ],\n')
        if entry.data.Tmin is not None: f.write('        Tmin = {0!r},\n'.format(entry.data.Tmin))
        if entry.data.Tmax is not None: f.write('        Tmax = {0!r},\n'.format(entry.data.Tmax))
        f.write('    ),\n')
    else:
        f.write('    thermo = {0!r},\n'.format(entry.data))

    if entry.reference is not None: f.write('    reference = {0!r},\n'.format(entry.reference))
    if entry.referenceType != "": f.write('    referenceType = "{0}",\n'.format(entry.referenceType))
    f.write('    shortDesc = u"""{0}""",\n'.format(entry.shortDesc))
    f.write('    longDesc = \n')
    f.write('u"""\n')
    f.write(entry.longDesc.strip() + "\n")
    f.write('""",\n')

    f.write('    history = [\n')
    for time, user, action, description in entry.history:
        f.write('        ("{0}","{1}","{2}","""{3}"""),\n'.format(time, user, action, description))
    f.write('    ],\n')

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

def generateOldLibraryEntry(data):
    """
    Return a list of values used to save entries to the old-style RMG
    thermo database based on the thermodynamics object `data`.
    """
    if isinstance(data, ThermoData):
        return '{0:9g} {1:9g} {2:9g} {3:9g} {4:9g} {5:9g} {6:9g} {7:9g} {8:9g} {9:9g} {10:9g} {11:9g}'.format(
            data.H298.value/4184.,
            data.S298.value/4.184,
            data.Cpdata.values[0]/4.184,
            data.Cpdata.values[1]/4.184,
            data.Cpdata.values[2]/4.184,
            data.Cpdata.values[3]/4.184,
            data.Cpdata.values[4]/4.184,
            data.Cpdata.values[5]/4.184,
            data.Cpdata.values[6]/4.184,
            data.H298.uncertainty/4184.,
            data.S298.uncertainty/4.184,
            data.Cpdata.uncertainty/4.184,
        )
    else:
        return '{0:9g} {1:9g} {2:9g} {3:9g} {4:9g} {5:9g} {6:9g} {7:9g} {8:9g} {9:9g} {10:9g} {11:9g}'.format(
            data.getEnthalpy(298)/4184.,
            data.getEntropy(298)/4.184,
            data.getHeatCapacity(300)/4.184,
            data.getHeatCapacity(400)/4.184,
            data.getHeatCapacity(500)/4.184,
            data.getHeatCapacity(600)/4.184,
            data.getHeatCapacity(800)/4.184,
            data.getHeatCapacity(1000)/4.184,
            data.getHeatCapacity(1500)/4.184,
            0,
            0,
            0,
        )

def processOldLibraryEntry(data):
    """
    Process a list of parameters `data` as read from an old-style RMG
    thermo database, returning the corresponding thermodynamics object.
    """
    return ThermoData(
        Tdata = ([300,400,500,600,800,1000,1500],"K"),
        Cpdata = ([float(d) for d in data[2:9]],"cal/(mol*K)","+|-",float(data[11])),
        H298 = (float(data[0]),"kcal/mol","+|-",float(data[9])),
        S298 = (float(data[1]),"cal/(mol*K)","+|-",float(data[10])),
    )

################################################################################

class ThermoDepository(Database):
    """
    A class for working with the RMG thermodynamics depository.
    """

    def __init__(self, label='', name='', shortDesc='', longDesc=''):
        Database.__init__(self, label=label, name=name, shortDesc=shortDesc, longDesc=longDesc)

    def loadEntry(self, index, label, molecule, thermo, reference=None, referenceType='', shortDesc='', longDesc='', history=None):
        self.entries[label] = Entry(
            index = index,
            label = label,
            item = Molecule().fromAdjacencyList(molecule),
            data = thermo,
            reference = reference,
            referenceType = referenceType,
            shortDesc = shortDesc,
            longDesc = longDesc.strip(),
            history = history or [],
        )

    def saveEntry(self, f, entry):
        """
        Write the given `entry` in the thermo database to the file object `f`.
        """
        return saveEntry(f, entry)

################################################################################

class ThermoLibrary(Database):
    """
    A class for working with a RMG thermodynamics library.
    """

    def __init__(self, label='', name='', shortDesc='', longDesc=''):
        Database.__init__(self, label=label, name=name, shortDesc=shortDesc, longDesc=longDesc)

    def loadEntry(self,
                  index,
                  label,
                  molecule,
                  thermo,
                  reference=None,
                  referenceType='',
                  shortDesc='',
                  longDesc='',
                  history=None
                  ):
        self.entries[label] = Entry(
            index = index,
            label = label,
            item = Molecule().fromAdjacencyList(molecule),
            data = thermo,
            reference = reference,
            referenceType = referenceType,
            shortDesc = shortDesc,
            longDesc = longDesc.strip(),
            history = history or [],
        )

    def saveEntry(self, f, entry):
        """
        Write the given `entry` in the thermo database to the file object `f`.
        """
        return saveEntry(f, entry)

    def generateOldLibraryEntry(self, data):
        """
        Return a list of values used to save entries to the old-style RMG
        thermo database based on the thermodynamics object `data`.
        """
        return generateOldLibraryEntry(data)

    def processOldLibraryEntry(self, data):
        """
        Process a list of parameters `data` as read from an old-style RMG
        thermo database, returning the corresponding thermodynamics object.
        """
        return processOldLibraryEntry(data)

################################################################################

class ThermoGroups(Database):
    """
    A class for working with an RMG thermodynamics group additivity database.
    """

    def __init__(self, label='', name='', shortDesc='', longDesc=''):
        Database.__init__(self, label=label, name=name, shortDesc=shortDesc, longDesc=longDesc)

    def loadEntry(self,
                  index,
                  label,
                  group,
                  thermo,
                  reference=None,
                  referenceType='',
                  shortDesc='',
                  longDesc='',
                  history=None
                  ):
        if group[0:3].upper() == 'OR{' or group[0:4].upper() == 'AND{' or group[0:7].upper() == 'NOT OR{' or group[0:8].upper() == 'NOT AND{':
            item = makeLogicNode(group)
        else:
            item = Group().fromAdjacencyList(group)
        self.entries[label] = Entry(
            index = index,
            label = label,
            item = item,
            data = thermo,
            reference = reference,
            referenceType = referenceType,
            shortDesc = shortDesc,
            longDesc = longDesc.strip(),
            history = history or [],
        )

    def saveEntry(self, f, entry):
        """
        Write the given `entry` in the thermo database to the file object `f`.
        """
        return saveEntry(f, entry)

    def generateOldLibraryEntry(self, data):
        """
        Return a list of values used to save entries to the old-style RMG
        thermo database based on the thermodynamics object `data`.
        """
        return generateOldLibraryEntry(data)

    def processOldLibraryEntry(self, data):
        """
        Process a list of parameters `data` as read from an old-style RMG
        thermo database, returning the corresponding thermodynamics object.
        """
        return processOldLibraryEntry(data)

################################################################################

class ThermoDatabase:
    """
    A class for working with the RMG thermodynamics database.
    """

    def __init__(self):
        self.depository = {}
        self.libraries = {}
        self.groups = {}
        self.libraryOrder = []
        self.local_context = {
            'ThermoData': ThermoData,
            'Wilhoit': Wilhoit,
            'NASA': NASA,
            'MultiNASA': MultiNASA,
        }
        self.global_context = {}

    def load(self, path, libraries=None, depository=True):
        """
        Load the thermo database from the given `path` on disk, where `path`
        points to the top-level folder of the thermo database.
        """
        if depository:
            self.loadDepository(os.path.join(path, 'depository'))
        else:
            self.depository = {}
        self.loadLibraries(os.path.join(path, 'libraries'), libraries)
        self.loadGroups(os.path.join(path, 'groups'))

    def loadDepository(self, path):
        """
        Load the thermo database from the given `path` on disk, where `path`
        points to the top-level folder of the thermo database.
        """
        self.depository = {}
        self.depository['stable']  = ThermoDepository().load(os.path.join(path, 'stable.py'), self.local_context, self.global_context)
        self.depository['radical'] = ThermoDepository().load(os.path.join(path, 'radical.py'), self.local_context, self.global_context)

    def loadLibraries(self, path, libraries=None):
        """
        Load the thermo database from the given `path` on disk, where `path`
        points to the top-level folder of the thermo database.
        """
        self.libraries = {}; self.libraryOrder = []
        for (root, dirs, files) in os.walk(os.path.join(path)):
            for f in files:
                name, ext = os.path.splitext(f)
                if ext.lower() == '.py' and (libraries is None or name in libraries):
                    logging.info('Loading thermodynamics library from {0} in {1}...'.format(f, root))
                    library = ThermoLibrary()
                    library.load(os.path.join(root, f), self.local_context, self.global_context)
                    library.label = os.path.splitext(f)[0]
                    self.libraries[library.label] = library
                    self.libraryOrder.append(library.label)
        if libraries is not None:
            self.libraryOrder = libraries

    def loadGroups(self, path):
        """
        Load the thermo database from the given `path` on disk, where `path`
        points to the top-level folder of the thermo database.
        """
        logging.info('Loading thermodynamics group database from {0}...'.format(path))
        self.groups = {}
        self.groups['group']   =   ThermoGroups(label='group').load(os.path.join(path, 'group.py'  ), self.local_context, self.global_context)
        self.groups['gauche']  =  ThermoGroups(label='gauche').load(os.path.join(path, 'gauche.py' ), self.local_context, self.global_context)
        self.groups['int15']   =   ThermoGroups(label='int15').load(os.path.join(path, 'int15.py'  ), self.local_context, self.global_context)
        self.groups['ring']    =    ThermoGroups(label='ring').load(os.path.join(path, 'ring.py'   ), self.local_context, self.global_context)
        self.groups['radical'] = ThermoGroups(label='radical').load(os.path.join(path, 'radical.py'), self.local_context, self.global_context)
        self.groups['other']   =   ThermoGroups(label='other').load(os.path.join(path, 'other.py'  ), self.local_context, self.global_context)

    def save(self, path):
        """
        Save the thermo database to the given `path` on disk, where `path`
        points to the top-level folder of the thermo database.
        """
        path = os.path.abspath(path)
        if not os.path.exists(path): os.mkdir(path)
        self.saveDepository(os.path.join(path, 'depository'))
        self.saveLibraries(os.path.join(path, 'libraries'))
        self.saveGroups(os.path.join(path, 'groups'))

    def saveDepository(self, path):
        """
        Save the thermo depository to the given `path` on disk, where `path`
        points to the top-level folder of the thermo depository.
        """
        if not os.path.exists(path): os.mkdir(path)
        self.depository['stable'].save(os.path.join(path, 'stable.py'))
        self.depository['radical'].save(os.path.join(path, 'radical.py'))

    def saveLibraries(self, path):
        """
        Save the thermo libraries to the given `path` on disk, where `path`
        points to the top-level folder of the thermo libraries.
        """
        if not os.path.exists(path): os.mkdir(path)
        for library in self.libraries.values():
            library.save(os.path.join(path, '{0}.py'.format(library.label)))

    def saveGroups(self, path):
        """
        Save the thermo groups to the given `path` on disk, where `path`
        points to the top-level folder of the thermo groups.
        """
        if not os.path.exists(path): os.mkdir(path)
        self.groups['group'].save(os.path.join(path, 'group.py'))
        self.groups['gauche'].save(os.path.join(path, 'gauche.py'))
        self.groups['int15'].save(os.path.join(path, 'int15.py'))
        self.groups['ring'].save(os.path.join(path, 'ring.py'))
        self.groups['radical'].save(os.path.join(path, 'radical.py'))
        self.groups['other'].save(os.path.join(path, 'other.py'))

    def loadOld(self, path):
        """
        Load the old RMG thermo database from the given `path` on disk, where
        `path` points to the top-level folder of the old RMG database.
        """
        # The old database does not have a depository, so create an empty one
        self.depository = {}
        self.depository['stable']  = ThermoDepository(label='stable', name='Stable Molecules')
        self.depository['radical'] = ThermoDepository(label='radical', name='Radical Molecules')

        for (root, dirs, files) in os.walk(os.path.join(path, 'thermo_libraries')):
            if os.path.exists(os.path.join(root, 'Dictionary.txt')) and os.path.exists(os.path.join(root, 'Library.txt')):
                library = ThermoLibrary(label=os.path.basename(root), name=os.path.basename(root))
                library.loadOld(
                    dictstr = os.path.join(root, 'Dictionary.txt'),
                    treestr = '',
                    libstr = os.path.join(root, 'Library.txt'),
                    numParameters = 12,
                    numLabels = 1,
                    pattern = False,
                )
                library.label = os.path.basename(root)
                self.libraries[library.label] = library

        self.groups = {}
        self.groups['group'] = ThermoGroups(label='group', name='Functional Group Additivity Values').loadOld(
            dictstr = os.path.join(path, 'thermo_groups', 'Group_Dictionary.txt'),
            treestr = os.path.join(path, 'thermo_groups', 'Group_Tree.txt'),
            libstr = os.path.join(path, 'thermo_groups', 'Group_Library.txt'),
            numParameters = 12,
            numLabels = 1,
            pattern = True,
        )
        self.groups['gauche'] = ThermoGroups(label='gauche', name='Gauche Interaction Corrections').loadOld(
            dictstr = os.path.join(path, 'thermo_groups', 'Gauche_Dictionary.txt'),
            treestr = os.path.join(path, 'thermo_groups', 'Gauche_Tree.txt'),
            libstr = os.path.join(path, 'thermo_groups', 'Gauche_Library.txt'),
            numParameters = 12,
            numLabels = 1,
            pattern = True,
        )
        self.groups['int15'] = ThermoGroups(label='int15', name='1,5-Interaction Corrections').loadOld(
            dictstr = os.path.join(path, 'thermo_groups', '15_Dictionary.txt'),
            treestr = os.path.join(path, 'thermo_groups', '15_Tree.txt'),
            libstr = os.path.join(path, 'thermo_groups', '15_Library.txt'),
            numParameters = 12,
            numLabels = 1,
            pattern = True,
        )
        self.groups['radical'] = ThermoGroups(label='radical', name='Radical Corrections').loadOld(
            dictstr = os.path.join(path, 'thermo_groups', 'Radical_Dictionary.txt'),
            treestr = os.path.join(path, 'thermo_groups', 'Radical_Tree.txt'),
            libstr = os.path.join(path, 'thermo_groups', 'Radical_Library.txt'),
            numParameters = 12,
            numLabels = 1,
            pattern = True,
        )
        self.groups['ring'] = ThermoGroups(label='ring', name='Ring Corrections').loadOld(
            dictstr = os.path.join(path, 'thermo_groups', 'Ring_Dictionary.txt'),
            treestr = os.path.join(path, 'thermo_groups', 'Ring_Tree.txt'),
            libstr = os.path.join(path, 'thermo_groups', 'Ring_Library.txt'),
            numParameters = 12,
            numLabels = 1,
            pattern = True,
        )
        self.groups['other'] = ThermoGroups(label='other', name='Other Corrections').loadOld(
            dictstr = os.path.join(path, 'thermo_groups', 'Other_Dictionary.txt'),
            treestr = os.path.join(path, 'thermo_groups', 'Other_Tree.txt'),
            libstr = os.path.join(path, 'thermo_groups', 'Other_Library.txt'),
            numParameters = 12,
            numLabels = 1,
            pattern = True,
        )

    def saveOld(self, path):
        """
        Save the old RMG thermo database to the given `path` on disk, where
        `path` points to the top-level folder of the old RMG database.
        """

        # Depository not used in old database, so it is not saved

        librariesPath = os.path.join(path, 'thermo_libraries')
        if not os.path.exists(librariesPath): os.mkdir(librariesPath)
        for library in self.libraries.values():
            libraryPath = os.path.join(librariesPath, library.label)
            if not os.path.exists(libraryPath): os.mkdir(libraryPath)
            library.saveOld(
                dictstr = os.path.join(libraryPath, 'Dictionary.txt'),
                treestr = '',
                libstr = os.path.join(libraryPath, 'Library.txt'),
            )

        groupsPath = os.path.join(path, 'thermo_groups')
        if not os.path.exists(groupsPath): os.mkdir(groupsPath)
        self.groups['group'].saveOld(
            dictstr = os.path.join(groupsPath, 'Group_Dictionary.txt'),
            treestr = os.path.join(groupsPath, 'Group_Tree.txt'),
            libstr = os.path.join(groupsPath, 'Group_Library.txt'),
        )
        self.groups['gauche'].saveOld(
            dictstr = os.path.join(groupsPath, 'Gauche_Dictionary.txt'),
            treestr = os.path.join(groupsPath, 'Gauche_Tree.txt'),
            libstr = os.path.join(groupsPath, 'Gauche_Library.txt'),
        )
        self.groups['int15'].saveOld(
            dictstr = os.path.join(groupsPath, '15_Dictionary.txt'),
            treestr = os.path.join(groupsPath, '15_Tree.txt'),
            libstr = os.path.join(groupsPath, '15_Library.txt'),
        )
        self.groups['radical'].saveOld(
            dictstr = os.path.join(groupsPath, 'Radical_Dictionary.txt'),
            treestr = os.path.join(groupsPath, 'Radical_Tree.txt'),
            libstr = os.path.join(groupsPath, 'Radical_Library.txt'),
        )
        self.groups['ring'].saveOld(
            dictstr = os.path.join(groupsPath, 'Ring_Dictionary.txt'),
            treestr = os.path.join(groupsPath, 'Ring_Tree.txt'),
            libstr = os.path.join(groupsPath, 'Ring_Library.txt'),
        )
        self.groups['other'].saveOld(
            dictstr = os.path.join(groupsPath, 'Other_Dictionary.txt'),
            treestr = os.path.join(groupsPath, 'Other_Tree.txt'),
            libstr = os.path.join(groupsPath, 'Other_Library.txt'),
        )

    def getThermoData(self, species):
        """
        Return the thermodynamic parameters for a given :class:`Species`
        object `species`. This function first searches the loaded libraries
        in order, returning the first match found, before falling back to
        estimation via group additivity.
        """
        thermoData = None
        # Check the libraries in order first; return the first successful match
        for label in self.libraryOrder:
            thermoData = self.getThermoDataFromLibrary(species, self.libraries[label])
            if thermoData is not None: break
        else:
            # Thermo not found in any loaded libraries, so estimate
            thermoData = self.getThermoDataFromGroups(species)
        return thermoData[0]

    def getAllThermoData(self, species):
        """
        Return all possible sets of thermodynamic parameters for a given
        :class:`Species` object `species`. The hits from the depository come
        first, then the libraries (in order), and then the group additivity
        estimate. This method is useful for a generic search job.
        """
        thermoData = []
        # Data from depository comes first
        thermoData.extend(self.getThermoDataFromDepository(species))
        # Data from libraries comes second
        for label in self.libraryOrder:
            data = self.getThermoDataFromLibrary(species, self.libraries[label])
            if data:
                thermoData.append(data)
        # Last entry is always the estimate from group additivity
        thermoData.append(self.getThermoDataFromGroups(species))
        return thermoData

    def getThermoDataFromDepository(self, species):
        """
        Return all possible sets of thermodynamic parameters for a given
        :class:`Species` object `species` from the depository. If no
        depository is loaded, a :class:`DatabaseError` is raised.
        """
        items = []
        for label, entry in self.depository['stable'].entries.iteritems():
            for molecule in species.molecule:
                if molecule.isIsomorphic(entry.item):
                    items.append((deepcopy(entry.data), self.depository['stable'], entry))
                    break
        for label, entry in self.depository['radical'].entries.iteritems():
            for molecule in species.molecule:
                if molecule.isIsomorphic(entry.item):
                    items.append((deepcopy(entry.data), self.depository['radical'], entry))
                    break
        return items

    def getThermoDataFromLibrary(self, species, library):
        """
        Return the set of thermodynamic parameters corresponding to a given
        :class:`Species` object `species` from the specified thermodynamics
        `library`. If `library` is a string, the list of libraries is searched
        for a library with that name. If no match is found in that library,
        ``None`` is returned. If no corresponding library is found, a
        :class:`DatabaseError` is raised.
        """
        for label, entry in library.entries.iteritems():
            for molecule in species.molecule:
                if molecule.isIsomorphic(entry.item) and entry.data is not None:
                    return (deepcopy(entry.data), library, entry)
        return None

    def getThermoDataFromGroups(self, species):
        """
        Return the set of thermodynamic parameters corresponding to a given
        :class:`Species` object `species` by estimation using the group
        additivity values. If no group additivity values are loaded, a
        :class:`DatabaseError` is raised.
        """
        thermo = []
        for molecule in species.molecule:
            molecule.clearLabeledAtoms()
            molecule.updateAtomTypes()
            tdata = self.estimateThermoViaGroupAdditivity(molecule)
            thermo.append(tdata)

        H298 = numpy.array([t.getEnthalpy(298.) for t in thermo])
        indices = H298.argsort()

        species.molecule = [species.molecule[ind] for ind in indices]

        return (thermo[indices[0]], None, None)

    def estimateThermoViaGroupAdditivity(self, molecule):
        """
        Return the set of thermodynamic parameters corresponding to a given
        :class:`Molecule` object `molecule` by estimation using the group
        additivity values. If no group additivity values are loaded, a
        :class:`DatabaseError` is raised.
        """
        # For thermo estimation we need the atoms to already be sorted because we
        # iterate over them; if the order changes during the iteration then we
        # will probably not visit the right atoms, and so will get the thermo wrong
        molecule.sortVertices()

        thermoData = None

        if sum([atom.radicalElectrons for atom in molecule.atoms]) > 0: # radical species

            # Make a copy of the structure so we don't change the original
            saturatedStruct = molecule.copy(deep=True)

            # Saturate structure by replacing all radicals with bonds to
            # hydrogen atoms
            added = {}
            for atom in saturatedStruct.atoms:
                for i in range(atom.radicalElectrons):
                    H = Atom('H')
                    bond = Bond(atom, H, 'S')
                    saturatedStruct.addAtom(H)
                    saturatedStruct.addBond(bond)
                    if atom not in added:
                        added[atom] = []
                    added[atom].append([H, bond])
                    atom.decrementRadical()

            # Update the atom types of the saturated structure (not sure why
            # this is necessary, because saturating with H shouldn't be
            # changing atom types, but it doesn't hurt anything and is not
            # very expensive, so will do it anyway)
            saturatedStruct.updateConnectivityValues()
            saturatedStruct.sortVertices()
            saturatedStruct.updateAtomTypes()

            # Get thermo estimate for saturated form of structure
            thermoData = self.estimateThermoViaGroupAdditivity(saturatedStruct)
            assert thermoData is not None, "Thermo data of saturated {0} of molecule {1} is None!".format(saturatedStruct, molecule)
            # Undo symmetry number correction for saturated structure
            thermoData.S298.value += constants.R * math.log(saturatedStruct.symmetryNumber)

            # For each radical site, get radical correction
            # Only one radical site should be considered at a time; all others
            # should be saturated with hydrogen atoms
            for atom in added:

                # Remove the added hydrogen atoms and bond and restore the radical
                for H, bond in added[atom]:
                    saturatedStruct.removeBond(bond)
                    saturatedStruct.removeAtom(H)
                    atom.incrementRadical()

                saturatedStruct.updateConnectivityValues()

                try:
                    thermoData = self.__addThermoData(thermoData, self.__getGroupThermoData(self.groups['radical'], saturatedStruct, {'*':atom}))
                except KeyError:
                    logging.error("Couldn't find in radical thermo database:")
                    logging.error(molecule)
                    logging.error(molecule.toAdjacencyList())
                    raise

                # Re-saturate
                for H, bond in added[atom]:
                    saturatedStruct.addAtom(H)
                    saturatedStruct.addBond(bond)
                    atom.decrementRadical()

                # Subtract the enthalpy of the added hydrogens
                for H, bond in added[atom]:
                    thermoData.H298.value -= 52.103 * 4184

            # Correct the entropy for the symmetry number

        else: # non-radical species
            # Generate estimate of thermodynamics
            for atom in molecule.atoms:
                # Iterate over heavy (non-hydrogen) atoms
                if atom.isNonHydrogen():
                    # Get initial thermo estimate from main group database
                    try:
                        if thermoData is None:
                            thermoData = self.__getGroupThermoData(self.groups['group'], molecule, {'*':atom})
                        else:
                            thermoData = self.__addThermoData(thermoData, self.__getGroupThermoData(self.groups['group'], molecule, {'*':atom}))
                    except KeyError:
                        logging.error("Couldn't find in main thermo database:")
                        logging.error(molecule)
                        logging.error(molecule.toAdjacencyList())
                        raise
                    # Correct for gauche and 1,5- interactions
                    try:
                        thermoData = self.__addThermoData(thermoData, self.__getGroupThermoData(self.groups['gauche'], molecule, {'*':atom}))
                    except KeyError: pass
                    try:
                        thermoData = self.__addThermoData(thermoData, self.__getGroupThermoData(self.groups['int15'], molecule, {'*':atom}))
                    except KeyError: pass
                    try:
                        thermoData = self.__addThermoData(thermoData, self.__getGroupThermoData(self.groups['other'], molecule, {'*':atom}))
                    except KeyError: pass

            # Do ring corrections separately because we only want to match
            # each ring one time; this doesn't work yet
            rings = molecule.getSmallestSetOfSmallestRings()
            for ring in rings:
                # Make a temporary structure containing only the atoms in the ring
                # NB. if any of the ring corrections depend on ligands not in the ring, they will not be found!
                ringStructure = Molecule()
                newAtoms = dict()
                for atom in ring:
                    newAtoms[atom] = atom.copy()
                    ringStructure.addAtom(newAtoms[atom]) # (addAtom deletes the atom's bonds)
                for atom1 in ring:
                    for atom2 in ring:
                        if molecule.hasBond(atom1, atom2):
                            ringStructure.addBond(Bond(newAtoms[atom1], newAtoms[atom2], atom1.bonds[atom2].order ))

                # Get thermo correction for this ring
                try:
                    thermoData = self.__addThermoData(thermoData, self.__getGroupThermoData(self.groups['ring'], ringStructure, {}))
                except KeyError:
                    logging.error("Couldn't find in ring database:")
                    logging.error(ringStructure)
                    logging.error(ringStructure.toAdjacencyList())
                    raise

        # Correct entropy for symmetry number
        molecule.calculateSymmetryNumber()
        thermoData.S298.value -= constants.R * math.log(molecule.symmetryNumber)

        return thermoData

    def __addThermoData(self, thermoData1, thermoData2):
        """
        Add two :class:`ThermoData` objects `thermoData1` and `thermoData2`
        together, returning their sum as a new :class:`ThermoData` object.
        """
        if len(thermoData1.Tdata.values) != len(thermoData2.Tdata.values) or any([T1 != T2 for T1, T2 in zip(thermoData1.Tdata.values, thermoData2.Tdata.values)]):
            raise ThermoError('Cannot add these ThermoData objects due to their having different temperature points.')
        new = ThermoData(
            Tdata = (thermoData1.Tdata.values, thermoData1.Tdata.units),
            Cpdata = (thermoData1.Cpdata.values + thermoData2.Cpdata.values, thermoData1.Tdata.units),
            H298 = (thermoData1.H298.value + thermoData2.H298.value, thermoData1.Tdata.units),
            S298 = (thermoData1.S298.value + thermoData2.S298.value, thermoData1.Tdata.units),
        )
        if thermoData1.comment == '': new.comment = thermoData2.comment
        elif thermoData2.comment == '': new.comment = thermoData1.comment
        else: new.comment = thermoData1.comment + ' + ' + thermoData2.comment
        return new

    def __getGroupThermoData(self, database, molecule, atom):
        """
        Determine the group additivity thermodynamic data for the atom `atom`
        in the structure `structure`.
        """

        node0 = database.descendTree(molecule, atom, None)

        if node0 is None:
            raise KeyError('Node not found in database.')

        # It's possible (and allowed) that items in the tree may not be in the
        # library, in which case we need to fall up the tree until we find an
        # ancestor that has an entry in the library
        node = node0
        while node.data is None and node is not None:
            node = node.parent
        if node is None:
            raise InvalidDatabaseError('Unable to determine thermo parameters for {0}: no library entries for {1} or any of its ancestors.'.format(molecule, node0) )

        data = node.data; comment = node.label
        while isinstance(data, str) and data is not None:
            for entry in database.entries.values():
                if entry.label == data:
                    data = entry.data
                    comment = entry.label
                    break
        data = deepcopy(data)
        data.comment = '{0}({1})'.format(database.label, comment)

        # This code prints the hierarchy of the found node; useful for debugging
        #result = ''
        #while node is not None:
        #   result = ' -> ' + node + result
        #   node = database.tree.parent[node]
        #print result[4:]

        return data


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import functools
import numpy as np
import pycuda.driver as cuda
from graphdot.cuda.array import umempty, umzeros
from graphdot.codegen import Template
from graphdot.codegen.cpptool import decltype
from graphdot.microkernel import TensorProduct, Product
from graphdot.kernel.marginalized._backend_cuda import CUDABackend
from graphdot.kernel.marginalized._octilegraph import OctileGraph

class AltCUDABackend(CUDABackend):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def allocate_pcg_scratch(self, number, max_graph_size, traits):
        length = max_graph_size**2
        n_temporaries = 5

        self.scratch_pcg, self.scratch_pcg_d = self._allocate_scratch(
            self.scratch_pcg, self.scratch_pcg_d, number, length,
            n_temporaries
        )
        return self.scratch_pcg_d

    @property
    @functools.lru_cache(maxsize=1)
    def template(self):
        return Template(os.path.join(os.path.dirname(__file__), 'template.cu'))

    def __call__(self, graphs, node_kernel, edge_kernel, p, q, eps, ftol, gtol,
                 jobs, gramian, traits, timer):
        ''' transfer graphs and starting probabilities to GPU '''
        timer.tic('transferring graphs to GPU')

        og_last = None
        graphs_d = umempty(len(graphs), dtype=OctileGraph.dtype)
        for i, g in enumerate(graphs):
            og, ogstate = self._register_graph(g)
            if i > 0:
                self._assert_homogeneous(og_last, og)
            og_last = og
            graphs_d[i] = ogstate

        weighted = og_last.weighted
        node_t = og_last.node_t
        edge_t = og_last.edge_t

        timer.toc('transferring graphs to GPU')

        ''' allocate global job counter '''
        timer.tic('allocate global job counter')
        i_job_global = umzeros(1, np.uint32)
        timer.toc('allocate global job counter')

        ''' code generation '''
        timer.tic('code generation')
        if weighted:
            edge_kernel = TensorProduct(weight=Product(),
                                        label=edge_kernel)

        node_kernel_src = self.gencode_kernel(node_kernel, 'node_kernel')
        edge_kernel_src = self.gencode_kernel(edge_kernel, 'edge_kernel')
        p_start_src = self.gencode_probability(p, 'p_start')

        with self.template.context(traits=traits) as template:
            self.source = template.render(
                node_kernel=node_kernel_src,
                edge_kernel=edge_kernel_src,
                p_start=p_start_src,
                node_t=decltype(node_t),
                edge_t=decltype(edge_t)
            )
        timer.toc('code generation')

        ''' JIT '''
        timer.tic('JIT')
        kernel = self.module.get_function('alt_graph_kernel_solver')
        timer.toc('JIT')

        ''' calculate launch configuration '''
        timer.tic('calculating launch configuration')
        launch_block_count = (self.device.MULTIPROCESSOR_COUNT
                              * self.block_per_sm)
        shmem_bytes_per_warp = self.module.get_global(
            'shmem_bytes_per_warp'
        )[1]
        shmem_bytes_per_block = (shmem_bytes_per_warp * self.block_size
                                 // self.device.WARP_SIZE)

        ''' allocate scratch buffers '''
        max_graph_size = np.max([len(g.nodes) for g in graphs])
        scratch_pcg = self.allocate_pcg_scratch(
            launch_block_count, max_graph_size, traits
        )

        ''' copy micro kernel parameters to GPU '''
        for name, uker in [('node_kernel', node_kernel),
                           ('edge_kernel', edge_kernel)]:
            states = np.array(
                self.pack_state(uker, diff_grid=False, diff_eps=eps),
                dtype=uker.dtype
            )

            p_uker, _ = self.module.get_global(name)
            cuda.memcpy_htod(p_uker, states[:1])

        p_p_start, _ = self.module.get_global('p_start')
        cuda.memcpy_htod(
            p_p_start, np.array([p.state], dtype=p.dtype)
        )

        timer.toc('calculating launch configuration')

        ''' GPU kernel execution '''
        timer.tic('GPU kernel execution')
        kernel(
            graphs_d,
            scratch_pcg,
            jobs,
            gramian,
            i_job_global,
            np.uint32(len(jobs)),
            np.float32(q),
            np.float32(q),  # placeholder for q0
            np.float32(eps),
            np.float32(ftol),
            np.float32(gtol),
            grid=(launch_block_count, 1, 1),
            block=(self.block_size, 1, 1),
            shared=shmem_bytes_per_block,
        )
        self.ctx.synchronize()
        timer.toc('GPU kernel execution')


# Author: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#         Daniel Strohmeier <daniel.strohmeier@gmail.com>
#
# License: Simplified BSD

import numpy as np
from scipy import linalg

from ..source_estimate import (SourceEstimate, VolSourceEstimate,
                               _BaseSourceEstimate)
from ..minimum_norm.inverse import (combine_xyz, _prepare_forward,
                                    _check_reference)
from ..forward import is_fixed_orient
from ..io.pick import pick_channels_evoked
from ..io.proj import deactivate_proj
from ..utils import logger, verbose, _check_depth
from ..dipole import Dipole

from .mxne_optim import (mixed_norm_solver, iterative_mixed_norm_solver, _Phi,
                         norm_l2inf, tf_mixed_norm_solver, norm_epsilon_inf)

@verbose
def _prepare_weights(forward, gain, source_weighting, weights, weights_min):
    mask = None
    if isinstance(weights, _BaseSourceEstimate):
        weights = np.max(np.abs(weights.data), axis=1)
    weights_max = np.max(weights)
    if weights_min > weights_max:
        raise ValueError('weights_min > weights_max (%s > %s)' %
                         (weights_min, weights_max))
    weights_min = weights_min / weights_max
    weights = weights / weights_max
    n_dip_per_pos = 1 if is_fixed_orient(forward) else 3
    weights = np.ravel(np.tile(weights, [n_dip_per_pos, 1]).T)
    if len(weights) != gain.shape[1]:
        raise ValueError('weights do not have the correct dimension '
                         ' (%d != %d)' % (len(weights), gain.shape[1]))
    if len(source_weighting.shape) == 1:
        source_weighting *= weights
    else:
        source_weighting *= weights[:, None]
    gain *= weights[None, :]

    if weights_min is not None:
        mask = (weights > weights_min)
        gain = gain[:, mask]
        n_sources = np.sum(mask) // n_dip_per_pos
        logger.info("Reducing source space to %d sources" % n_sources)

    return gain, source_weighting, mask

def _prepare_gain(forward, info, noise_cov, pca, depth, loose, rank,
                  weights=None, weights_min=None):
    depth = _check_depth(depth, 'depth_sparse')
    forward, gain_info, gain, _, _, source_weighting, _, _, whitener = \
        _prepare_forward(forward, info, noise_cov, 'auto', loose, rank, pca,
                         use_cps=True, **depth)

    if weights is None:
        mask = None
    else:
        gain, source_weighting, mask = _prepare_weights(
            forward, gain, source_weighting, weights, weights_min)

    return forward, gain, gain_info, whitener, source_weighting, mask

def _reapply_source_weighting(X, source_weighting, active_set):
    X *= source_weighting[active_set][:, None]
    return X

def _compute_residual(forward, evoked, X, active_set, info):
    # OK, picking based on row_names is safe
    sel = [forward['sol']['row_names'].index(c) for c in info['ch_names']]
    residual = evoked.copy()
    residual = pick_channels_evoked(residual, include=info['ch_names'])
    r_tmp = residual.copy()

    r_tmp.data = np.dot(forward['sol']['data'][sel, :][:, active_set], X)

    # Take care of proj
    active_projs = list()
    non_active_projs = list()
    for p in evoked.info['projs']:
        if p['active']:
            active_projs.append(p)
        else:
            non_active_projs.append(p)

    if len(active_projs) > 0:
        r_tmp.info['projs'] = deactivate_proj(active_projs, copy=True)
        r_tmp.apply_proj()
        r_tmp.add_proj(non_active_projs, remove_existing=False)

    residual.data -= r_tmp.data

    return residual

@verbose
def _make_sparse_stc(X, active_set, forward, tmin, tstep,
                     active_is_idx=False, verbose=None):
    if not is_fixed_orient(forward):
        logger.info('combining the current components...')
        X = combine_xyz(X)

    if not active_is_idx:
        active_idx = np.where(active_set)[0]
    else:
        active_idx = active_set

    n_dip_per_pos = 1 if is_fixed_orient(forward) else 3
    if n_dip_per_pos > 1:
        active_idx = np.unique(active_idx // n_dip_per_pos)

    src = forward['src']

    if src.kind != 'surface':
        vertices = src[0]['vertno'][active_idx]
        stc = VolSourceEstimate(X, vertices=vertices, tmin=tmin, tstep=tstep)
    else:
        vertices = []
        n_points_so_far = 0
        for this_src in src:
            this_n_points_so_far = n_points_so_far + len(this_src['vertno'])
            this_active_idx = active_idx[(n_points_so_far <= active_idx) &
                                         (active_idx < this_n_points_so_far)]
            this_active_idx -= n_points_so_far
            this_vertno = this_src['vertno'][this_active_idx]
            n_points_so_far = this_n_points_so_far
            vertices.append(this_vertno)

        stc = SourceEstimate(X, vertices=vertices, tmin=tmin, tstep=tstep)

    return stc

@verbose
def _make_dipoles_sparse(X, active_set, forward, tmin, tstep, M, M_est,
                         active_is_idx=False, verbose=None):
    times = tmin + tstep * np.arange(X.shape[1])

    if not active_is_idx:
        active_idx = np.where(active_set)[0]
    else:
        active_idx = active_set

    n_dip_per_pos = 1 if is_fixed_orient(forward) else 3
    if n_dip_per_pos > 1:
        active_idx = np.unique(active_idx // n_dip_per_pos)

    gof = np.zeros(M_est.shape[1])
    M_norm2 = np.sum(M ** 2, axis=0)
    R_norm2 = np.sum((M - M_est) ** 2, axis=0)
    gof[M_norm2 > 0.0] = 1. - R_norm2[M_norm2 > 0.0] / M_norm2[M_norm2 > 0.0]
    gof *= 100.

    dipoles = []
    for k, i_dip in enumerate(active_idx):
        i_pos = forward['source_rr'][i_dip][np.newaxis, :]
        i_pos = i_pos.repeat(len(times), axis=0)
        X_ = X[k * n_dip_per_pos: (k + 1) * n_dip_per_pos]
        if n_dip_per_pos == 1:
            amplitude = X_[0]
            i_ori = forward['source_nn'][i_dip][np.newaxis, :]
            i_ori = i_ori.repeat(len(times), axis=0)
        else:
            if forward['surf_ori']:
                X_ = np.dot(forward['source_nn'][
                    i_dip * n_dip_per_pos:(i_dip + 1) * n_dip_per_pos].T, X_)

            amplitude = np.sqrt(np.sum(X_ ** 2, axis=0))
            i_ori = np.zeros((len(times), 3))
            i_ori[amplitude > 0.] = (X_[:, amplitude > 0.] /
                                     amplitude[amplitude > 0.]).T

        dipoles.append(Dipole(times, i_pos, amplitude, i_ori, gof))

    return dipoles

@verbose
def make_stc_from_dipoles(dipoles, src, verbose=None):
    """Convert a list of spatio-temporal dipoles into a SourceEstimate.

    Parameters
    ----------
    dipoles : Dipole | list of instances of Dipole
        The dipoles to convert.
    src : instance of SourceSpaces
        The source space used to generate the forward operator.
    %(verbose)s

    Returns
    -------
    stc : SourceEstimate
        The source estimate.
    """
    logger.info('Converting dipoles into a SourceEstimate.')
    if isinstance(dipoles, Dipole):
        dipoles = [dipoles]
    if not isinstance(dipoles, list):
        raise ValueError('Dipoles must be an instance of Dipole or '
                         'a list of instances of Dipole. '
                         'Got %s!' % type(dipoles))
    tmin = dipoles[0].times[0]
    tstep = dipoles[0].times[1] - tmin
    X = np.zeros((len(dipoles), len(dipoles[0].times)))
    source_rr = np.concatenate([_src['rr'][_src['vertno'], :] for _src in src],
                               axis=0)
    n_lh_points = len(src[0]['vertno'])
    lh_vertno = list()
    rh_vertno = list()
    for i in range(len(dipoles)):
        if not np.all(dipoles[i].pos == dipoles[i].pos[0]):
            raise ValueError('Only dipoles with fixed position over time '
                             'are supported!')
        X[i] = dipoles[i].amplitude
        idx = np.all(source_rr == dipoles[i].pos[0], axis=1)
        idx = np.where(idx)[0][0]
        if idx < n_lh_points:
            lh_vertno.append(src[0]['vertno'][idx])
        else:
            rh_vertno.append(src[1]['vertno'][idx - n_lh_points])
    vertices = [np.array(lh_vertno).astype(int),
                np.array(rh_vertno).astype(int)]
    stc = SourceEstimate(X, vertices=vertices, tmin=tmin, tstep=tstep,
                         subject=src._subject)
    logger.info('[done]')
    return stc

@verbose
def mixed_norm(evoked, forward, noise_cov, alpha, loose='auto', depth=0.8,
               maxit=3000, tol=1e-4, active_set_size=10,
               debias=True, time_pca=True, weights=None, weights_min=0.,
               solver='auto', n_mxne_iter=1, return_residual=False,
               return_as_dipoles=False, dgap_freq=10, rank=None,
               verbose=None):
    """Mixed-norm estimate (MxNE) and iterative reweighted MxNE (irMxNE).

    Compute L1/L2 mixed-norm solution [1]_ or L0.5/L2 [2]_ mixed-norm
    solution on evoked data.

    Parameters
    ----------
    evoked : instance of Evoked or list of instances of Evoked
        Evoked data to invert.
    forward : dict
        Forward operator.
    noise_cov : instance of Covariance
        Noise covariance to compute whitener.
    alpha : float in range [0, 100)
        Regularization parameter. 0 means no regularization, 100 would give 0
        active dipole.
    loose : float in [0, 1] | 'auto'
        Value that weights the source variances of the dipole components
        that are parallel (tangential) to the cortical surface. If loose
        is 0 then the solution is computed with fixed orientation.
        If loose is 1, it corresponds to free orientations.
        The default value ('auto') is set to 0.2 for surface-oriented source
        space and set to 1.0 for volumic or discrete source space.
    %(depth)s
    maxit : int
        Maximum number of iterations.
    tol : float
        Tolerance parameter.
    active_set_size : int | None
        Size of active set increment. If None, no active set strategy is used.
    debias : bool
        Remove coefficient amplitude bias due to L1 penalty.
    time_pca : bool or int
        If True the rank of the concatenated epochs is reduced to
        its true dimension. If is 'int' the rank is limited to this value.
    weights : None | array | SourceEstimate
        Weight for penalty in mixed_norm. Can be None, a
        1d array with shape (n_sources,), or a SourceEstimate (e.g. obtained
        with wMNE, dSPM, or fMRI).
    weights_min : float
        Do not consider in the estimation sources for which weights
        is less than weights_min.
    solver : 'prox' | 'cd' | 'bcd' | 'auto'
        The algorithm to use for the optimization. 'prox' stands for
        proximal iterations using the FISTA algorithm, 'cd' uses
        coordinate descent, and 'bcd' applies block coordinate descent.
        'cd' is only available for fixed orientation.
    n_mxne_iter : int
        The number of MxNE iterations. If > 1, iterative reweighting
        is applied.
    return_residual : bool
        If True, the residual is returned as an Evoked instance.
    return_as_dipoles : bool
        If True, the sources are returned as a list of Dipole instances.
    dgap_freq : int or np.inf
        The duality gap is evaluated every dgap_freq iterations. Ignored if
        solver is 'cd'.
    %(rank_None)s

        .. versionadded:: 0.18
    %(verbose)s

    Returns
    -------
    stc : SourceEstimate | list of SourceEstimate
        Source time courses for each evoked data passed as input.
    residual : instance of Evoked
        The residual a.k.a. data not explained by the sources.
        Only returned if return_residual is True.

    See Also
    --------
    tf_mixed_norm

    References
    ----------
    .. [1] A. Gramfort, M. Kowalski, M. Hämäläinen,
       "Mixed-norm estimates for the M/EEG inverse problem using accelerated
       gradient methods", Physics in Medicine and Biology, 2012.
       https://doi.org/10.1088/0031-9155/57/7/1937

    .. [2] D. Strohmeier, Y. Bekhti, J. Haueisen, A. Gramfort,
       "The Iterative Reweighted Mixed-Norm Estimate for Spatio-Temporal
       MEG/EEG Source Reconstruction", IEEE Transactions of Medical Imaging,
       Volume 35 (10), pp. 2218-2228, 2016.
    """
    if not (0. <= alpha < 100.):
        raise ValueError('alpha must be in [0, 100). '
                         'Got alpha = %s' % alpha)
    if n_mxne_iter < 1:
        raise ValueError('MxNE has to be computed at least 1 time. '
                         'Requires n_mxne_iter >= 1, got %d' % n_mxne_iter)
    if dgap_freq <= 0.:
        raise ValueError('dgap_freq must be a positive integer.'
                         ' Got dgap_freq = %s' % dgap_freq)

    pca = True
    if not isinstance(evoked, list):
        evoked = [evoked]

    _check_reference(evoked[0])

    all_ch_names = evoked[0].ch_names
    if not all(all_ch_names == evoked[i].ch_names
               for i in range(1, len(evoked))):
        raise Exception('All the datasets must have the same good channels.')

    forward, gain, gain_info, whitener, source_weighting, mask = _prepare_gain(
        forward, evoked[0].info, noise_cov, pca, depth, loose, rank,
        weights, weights_min)

    sel = [all_ch_names.index(name) for name in gain_info['ch_names']]
    M = np.concatenate([e.data[sel] for e in evoked], axis=1)

    # Whiten data
    logger.info('Whitening data matrix.')
    M = np.dot(whitener, M)

    if time_pca:
        U, s, Vh = linalg.svd(M, full_matrices=False)
        if not isinstance(time_pca, bool) and isinstance(time_pca, int):
            U = U[:, :time_pca]
            s = s[:time_pca]
            Vh = Vh[:time_pca]
        M = U * s

    # Scaling to make setting of alpha easy
    n_dip_per_pos = 1 if is_fixed_orient(forward) else 3
    alpha_max = norm_l2inf(np.dot(gain.T, M), n_dip_per_pos, copy=False)
    alpha_max *= 0.01
    gain /= alpha_max
    source_weighting /= alpha_max

    if n_mxne_iter == 1:
        X, active_set, E = mixed_norm_solver(
            M, gain, alpha, maxit=maxit, tol=tol,
            active_set_size=active_set_size, n_orient=n_dip_per_pos,
            debias=debias, solver=solver, dgap_freq=dgap_freq, verbose=verbose)
    else:
        X, active_set, E = iterative_mixed_norm_solver(
            M, gain, alpha, n_mxne_iter, maxit=maxit, tol=tol,
            n_orient=n_dip_per_pos, active_set_size=active_set_size,
            debias=debias, solver=solver, dgap_freq=dgap_freq, verbose=verbose)

    if time_pca:
        X = np.dot(X, Vh)
        M = np.dot(M, Vh)

    # Compute estimated whitened sensor data
    M_estimated = np.dot(gain[:, active_set], X)

    if mask is not None:
        active_set_tmp = np.zeros(len(mask), dtype=np.bool)
        active_set_tmp[mask] = active_set
        active_set = active_set_tmp
        del active_set_tmp

    if active_set.sum() == 0:
        raise Exception("No active dipoles found. alpha is too big.")

    # Reapply weights to have correct unit
    X = _reapply_source_weighting(X, source_weighting, active_set)

    outs = list()
    residual = list()
    cnt = 0
    for e in evoked:
        tmin = e.times[0]
        tstep = 1.0 / e.info['sfreq']
        Xe = X[:, cnt:(cnt + len(e.times))]
        if return_as_dipoles:
            out = _make_dipoles_sparse(
                Xe, active_set, forward, tmin, tstep,
                M[:, cnt:(cnt + len(e.times))],
                M_estimated[:, cnt:(cnt + len(e.times))], verbose=None)
        else:
            out = _make_sparse_stc(Xe, active_set, forward, tmin, tstep)
        outs.append(out)
        cnt += len(e.times)

        if return_residual:
            residual.append(_compute_residual(forward, e, Xe, active_set,
                                              gain_info))

    logger.info('[done]')

    if len(outs) == 1:
        out = outs[0]
        if return_residual:
            residual = residual[0]
    else:
        out = outs

    if return_residual:
        out = out, residual

    return out

def _window_evoked(evoked, size):
    """Window evoked (size in seconds)."""
    if isinstance(size, (float, int)):
        lsize = rsize = float(size)
    else:
        lsize, rsize = size
    evoked = evoked.copy()
    sfreq = float(evoked.info['sfreq'])
    lsize = int(lsize * sfreq)
    rsize = int(rsize * sfreq)
    lhann = np.hanning(lsize * 2)[:lsize]
    rhann = np.hanning(rsize * 2)[-rsize:]
    window = np.r_[lhann, np.ones(len(evoked.times) - lsize - rsize), rhann]
    evoked.data *= window[None, :]
    return evoked

@verbose
def tf_mixed_norm(evoked, forward, noise_cov,
                  loose='auto', depth=0.8, maxit=3000,
                  tol=1e-4, weights=None, weights_min=0., pca=True,
                  debias=True, wsize=64, tstep=4, window=0.02,
                  return_residual=False, return_as_dipoles=False,
                  alpha=None, l1_ratio=None, dgap_freq=10, rank=None,
                  verbose=None):
    """Time-Frequency Mixed-norm estimate (TF-MxNE).

    Compute L1/L2 + L1 mixed-norm solution on time-frequency
    dictionary. Works with evoked data [1]_ [2]_.

    Parameters
    ----------
    evoked : instance of Evoked
        Evoked data to invert.
    forward : dict
        Forward operator.
    noise_cov : instance of Covariance
        Noise covariance to compute whitener.
    loose : float in [0, 1] | 'auto'
        Value that weights the source variances of the dipole components
        that are parallel (tangential) to the cortical surface. If loose
        is 0 then the solution is computed with fixed orientation.
        If loose is 1, it corresponds to free orientations.
        The default value ('auto') is set to 0.2 for surface-oriented source
        space and set to 1.0 for volumic or discrete source space.
    %(depth)s
    maxit : int
        Maximum number of iterations.
    tol : float
        Tolerance parameter.
    weights: None | array | SourceEstimate
        Weight for penalty in mixed_norm. Can be None or
        1d array of length n_sources or a SourceEstimate e.g. obtained
        with wMNE or dSPM or fMRI.
    weights_min: float
        Do not consider in the estimation sources for which weights
        is less than weights_min.
    pca: bool
        If True the rank of the data is reduced to true dimension.
    debias: bool
        Remove coefficient amplitude bias due to L1 penalty.
    wsize: int or array-like
        Length of the STFT window in samples (must be a multiple of 4).
        If an array is passed, multiple TF dictionaries are used (each having
        its own wsize and tstep) and each entry of wsize must be a multiple
        of 4. See [3]_.
    tstep: int or array-like
        Step between successive windows in samples (must be a multiple of 2,
        a divider of wsize and smaller than wsize/2) (default: wsize/2).
        If an array is passed, multiple TF dictionaries are used (each having
        its own wsize and tstep), and each entry of tstep must be a multiple
        of 2 and divide the corresponding entry of wsize. See [3]_.
    window : float or (float, float)
        Length of time window used to take care of edge artifacts in seconds.
        It can be one float or float if the values are different for left
        and right window length.
    return_residual : bool
        If True, the residual is returned as an Evoked instance.
    return_as_dipoles : bool
        If True, the sources are returned as a list of Dipole instances.
    alpha : float in [0, 100) or None
        Overall regularization parameter.
        If alpha and l1_ratio are not None, alpha_space and alpha_time are
        overridden by alpha * alpha_max * (1. - l1_ratio) and alpha * alpha_max
        * l1_ratio. 0 means no regularization, 100 would give 0 active dipole.
    l1_ratio : float in [0, 1] or None
        Proportion of temporal regularization.
        If l1_ratio and alpha are not None, alpha_space and alpha_time are
        overridden by alpha * alpha_max * (1. - l1_ratio) and alpha * alpha_max
        * l1_ratio. 0 means no time regularization aka MxNE.
    dgap_freq : int or np.inf
        The duality gap is evaluated every dgap_freq iterations.
    %(rank_None)s

        .. versionadded:: 0.18
    %(verbose)s

    Returns
    -------
    stc : instance of SourceEstimate
        Source time courses.
    residual : instance of Evoked
        The residual a.k.a. data not explained by the sources.
        Only returned if return_residual is True.

    See Also
    --------
    mixed_norm

    References
    ----------
    .. [1] A. Gramfort, D. Strohmeier, J. Haueisen, M. Hämäläinen, M. Kowalski
       "Time-Frequency Mixed-Norm Estimates: Sparse M/EEG imaging with
       non-stationary source activations",
       Neuroimage, Volume 70, pp. 410-422, 15 April 2013.
       DOI: 10.1016/j.neuroimage.2012.12.051

    .. [2] A. Gramfort, D. Strohmeier, J. Haueisen, M. Hämäläinen, M. Kowalski
       "Functional Brain Imaging with M/EEG Using Structured Sparsity in
       Time-Frequency Dictionaries",
       Proceedings Information Processing in Medical Imaging
       Lecture Notes in Computer Science, Volume 6801/2011, pp. 600-611, 2011.
       DOI: 10.1007/978-3-642-22092-0_49

    .. [3] Y. Bekhti, D. Strohmeier, M. Jas, R. Badeau, A. Gramfort.
       "M/EEG source localization with multiscale time-frequency dictionaries",
       6th International Workshop on Pattern Recognition in Neuroimaging
       (PRNI), 2016.
       DOI: 10.1109/PRNI.2016.7552337
    """
    _check_reference(evoked)

    all_ch_names = evoked.ch_names
    info = evoked.info

    if not (0. <= alpha < 100.):
        raise ValueError('alpha must be in [0, 100). '
                         'Got alpha = %s' % alpha)

    if not (0. <= l1_ratio <= 1.):
        raise ValueError('l1_ratio must be in range [0, 1].'
                         ' Got l1_ratio = %s' % l1_ratio)
    alpha_space = alpha * (1. - l1_ratio)
    alpha_time = alpha * l1_ratio

    if dgap_freq <= 0.:
        raise ValueError('dgap_freq must be a positive integer.'
                         ' Got dgap_freq = %s' % dgap_freq)

    tstep = np.atleast_1d(tstep)
    wsize = np.atleast_1d(wsize)
    if len(tstep) != len(wsize):
        raise ValueError('The same number of window sizes and steps must be '
                         'passed. Got tstep = %s and wsize = %s' %
                         (tstep, wsize))

    forward, gain, gain_info, whitener, source_weighting, mask = _prepare_gain(
        forward, evoked.info, noise_cov, pca, depth, loose, rank,
        weights, weights_min)
    n_dip_per_pos = 1 if is_fixed_orient(forward) else 3

    if window is not None:
        evoked = _window_evoked(evoked, window)

    sel = [all_ch_names.index(name) for name in gain_info["ch_names"]]
    M = evoked.data[sel]

    # Whiten data
    logger.info('Whitening data matrix.')
    M = np.dot(whitener, M)

    # Scaling to make setting of alpha easy
    n_steps = np.ceil(M.shape[1] / tstep.astype(float)).astype(int)
    n_freqs = wsize // 2 + 1
    n_coefs = n_steps * n_freqs
    phi = _Phi(wsize, tstep, n_coefs)

    alpha_max = norm_epsilon_inf(gain, M, phi, l1_ratio, n_dip_per_pos)
    alpha_max *= 0.01
    gain /= alpha_max
    source_weighting /= alpha_max

    X, active_set, E = tf_mixed_norm_solver(
        M, gain, alpha_space, alpha_time, wsize=wsize, tstep=tstep,
        maxit=maxit, tol=tol, verbose=verbose, n_orient=n_dip_per_pos,
        dgap_freq=dgap_freq, debias=debias)

    if active_set.sum() == 0:
        raise Exception("No active dipoles found. "
                        "alpha_space/alpha_time are too big.")

    # Compute estimated whitened sensor data
    M_estimated = np.dot(gain[:, active_set], X)

    if mask is not None:
        active_set_tmp = np.zeros(len(mask), dtype=np.bool)
        active_set_tmp[mask] = active_set
        active_set = active_set_tmp
        del active_set_tmp

    X = _reapply_source_weighting(X, source_weighting, active_set)

    if return_residual:
        residual = _compute_residual(
            forward, evoked, X, active_set, gain_info)

    if return_as_dipoles:
        out = _make_dipoles_sparse(
            X, active_set, forward, evoked.times[0], 1.0 / info['sfreq'],
            M, M_estimated, verbose=None)
    else:
        out = _make_sparse_stc(
            X, active_set, forward, evoked.times[0], 1.0 / info['sfreq'])

    logger.info('[done]')

    if return_residual:
        out = out, residual

    return out


import logging

from abc import ABC, abstractmethod

import numpy as np

from typing import Dict, Tuple, Any

from scipy.constants import epsilon_0
from scipy.integrate import trapz

from amset.misc.constants import k_B, e, hbar
from amset.data import AmsetData
from amset.misc.log import log_list
from amset.misc.util import f0
from pymatgen import Spin

__author__ = "Alex Ganose"
__maintainer__ = "Alex Ganose"
__email__ = "aganose@lbl.gov"
__date__ = "June 21, 2019"

logger = logging.getLogger(__name__)

class AbstractElasticScattering(ABC):

    name: str
    required_properties: Tuple[str]

    def __init__(self,
                 materials_properties: Dict[str, Any],
                 amset_data: AmsetData):
        self.properties = {p: materials_properties[p]
                           for p in self.required_properties}
        self.doping = amset_data.doping
        self.temperatures = amset_data.temperatures
        self.nbands = {s: len(amset_data.energies[s]) for s in amset_data.spins}
        self.spins = amset_data.spins

    @abstractmethod
    def prefactor(self, spin: Spin, b_idx: int):
        pass

    @abstractmethod
    def factor(self, k_diff_sq: np.ndarray):
        pass

class AcousticDeformationPotentialScattering(AbstractElasticScattering):

    name = "ACD"
    required_properties = ("deformation_potential", "elastic_constant")

    def __init__(self,
                 materials_properties: Dict[str, Any],
                 amset_data: AmsetData):
        super().__init__(materials_properties, amset_data)
        self.vb_idx = amset_data.vb_idx
        self.is_metal = amset_data.is_metal
        self._prefactor = (1e18 * e * k_B / (
                4.0 * np.pi ** 2 * hbar * self.properties["elastic_constant"]))

        self.deformation_potential = self.properties["deformation_potential"]
        if self.is_metal and isinstance(self.deformation_potential, tuple):
            logger.warning(
                "System is metallic but deformation potentials for both "
                "the valence and conduction bands have been set... using the "
                "valence band potential for all bands")
            self.deformation_potential = self.deformation_potential[0]

        elif not self.is_metal and not isinstance(
                self.deformation_potential, tuple):
            logger.warning(
                "System is semiconducting but only one deformation "
                "potential has been set... using this potential for all bands.")
            self.deformation_potential = (self.deformation_potential,
                                          self.deformation_potential)

    def prefactor(self, spin: Spin, b_idx: int):
        prefactor = self._prefactor * self.temperatures[None, :] * np.ones(
            (len(self.doping), len(self.temperatures)))

        if self.is_metal:
            prefactor *= self.properties["deformation_potential"] ** 2

        else:
            def_idx = 1 if b_idx > self.vb_idx[spin] else 0
            prefactor *= self.properties["deformation_potential"][def_idx] ** 2

        return prefactor

    def factor(self, k_diff_sq: np.ndarray):
        return np.ones((len(self.doping), len(self.temperatures),
                        k_diff_sq.shape[0]))

class IonizedImpurityScattering(AbstractElasticScattering):

    name = "IMP"
    required_properties = ("acceptor_charge", "donor_charge",
                           "static_dielectric")

    def __init__(self,
                 materials_properties: Dict[str, Any],
                 amset_data: AmsetData):
        super().__init__(materials_properties, amset_data)
        logger.debug("Initializing IMP scattering")

        self.beta_sq = np.zeros(amset_data.fermi_levels.shape)
        self.impurity_concentration = np.zeros(amset_data.fermi_levels.shape)

        tdos = amset_data.dos.tdos
        energies = amset_data.dos.energies
        fermi_levels = amset_data.fermi_levels
        vol = amset_data.structure.volume

        imp_info = []
        for n, t in np.ndindex(self.beta_sq.shape):
            ef = fermi_levels[n, t]
            temp = amset_data.temperatures[t]
            f = f0(energies, ef, temp)
            integral = trapz(tdos * f * (1 - f), x=energies)
            self.beta_sq[n, t] = (
                e ** 2 * integral * 1e12 /
                (self.properties["static_dielectric"] * epsilon_0 * k_B *
                 temp * e * vol))

            n_conc = np.abs(amset_data.electron_conc[n, t])
            p_conc = np.abs(amset_data.hole_conc[n, t])

            self.impurity_concentration[n, t] = (
                    n_conc * self.properties["donor_charge"] ** 2 +
                    p_conc * self.properties["acceptor_charge"] ** 2)
            imp_info.append(
                "{:.2g} cm⁻³ & {} K: β² = {:.4g}, Nᵢᵢ = {:.4g}".format(
                    amset_data.doping[n], temp, self.beta_sq[n, t],
                    self.impurity_concentration[n, t]))

        logger.debug("Inverse screening length (β) and impurity concentration "
                     "(Nᵢᵢ):")
        log_list(imp_info, level=logging.DEBUG)

        self._prefactor = (
                (1e-3 / (e ** 2)) * e ** 4 * self.impurity_concentration /
                (4.0 * np.pi ** 2 * epsilon_0 ** 2 * hbar *
                 self.properties["static_dielectric"] ** 2))

    def prefactor(self, spin: Spin, b_idx: int):
        # need to return prefactor with shape (nspins, ndops, ntemps, nbands)
        return self._prefactor

    def factor(self, k_diff_sq: np.ndarray):
        # tile k_diff_sq to make it commensurate with the dimensions of beta
        return 1 / (np.tile(k_diff_sq, (len(self.doping),
                                        len(self.temperatures), 1)) +
                    self.beta_sq[..., None]) ** 2

class PiezoelectricScattering(AbstractElasticScattering):

    name = "PIE"
    required_properties = ("piezoelectric_coefficient", "static_dielectric")

    def __init__(self,
                 materials_properties: Dict[str, Any],
                 amset_data: AmsetData):
        super().__init__(materials_properties, amset_data)
        unit_conversion = 1e9 / e
        self._prefactor = (unit_conversion * e ** 2 * k_B *
                           self.properties["piezoelectric_coefficient"] ** 2 /
                           (4.0 * np.pi ** 2 * hbar * epsilon_0 *
                            self.properties["static_dielectric"]))

    def prefactor(self, spin: Spin, b_idx: int):
        # need to return prefactor with shape (ndops, ntemps)
        return self._prefactor * self.temperatures[None, :] * np.ones(
                (len(self.doping), len(self.temperatures)))

    def factor(self, k_diff_sq: np.ndarray):
        # factor should have shape (ndops, ntemps, nkpts)
        return 1 / np.tile(k_diff_sq, (len(self.doping),
                                       len(self.temperatures), 1))