import random
import copy
import numpy as np
import pandas as pd
import os

from models.MultiGMPmodel import MultiCopyGMPmodel
from models.MultiGGHPmodel import MultiCopyGGHPmodel
from models.GGHPmodel import GGHPmodel
from util.calculateFissionsAndFusions import calculateFissionAndFussions

def readSequence(file):
    chr = []
    with open(file,'r') as rf:
        while True:
            line = rf.readline()[:-2]
            if not line:
                break
            itemset = line.split(' ')
            header = itemset[0]
            new_itemset = [header]
            for i in itemset[1:]:
                item = i.split('_')
                new_itemset.append(item[0])
            chr.append(new_itemset)
    return chr

def outSequence(sequence,outfile):
    outfile = open(outfile,'w')
    for i in sequence:
        for j in i:
            outfile.write(j+' ')
        outfile.write('\n')
    outfile.close()

def sequence2adjacency(sequence):
    adjacency = []
    for i in sequence:
        block = i[0]
        if block.startswith('-'):
            adjacency.append(['$',block[1:] + 'b'])
            start = block[1:] + 'a'
        else:
            adjacency.append(['$', block + 'a'])
            start = block + 'b'
        for j in range(len(i)-1):
            block = i[j+1]
            if block.startswith('-'):
                adjacency.append([start, block[1:] + 'b'])
                start = block[1:] + 'a'
            else:
                adjacency.append([start, block + 'a'])
                start = block + 'b'
        adjacency.append([start,'$'])
    return adjacency

def reverse(item):
    if item[-1] == 'a':
        return item[:-1] + 'b'
    else:
        return item[:-1] + 'a'

def assemble(adjacency_list):
    # build adj M
    matrix_items = []
    for i in adjacency_list:
        for j in i:
            if j not in matrix_items:
                matrix_items.append(j)
    matrix_items = sorted(matrix_items)
    adjacency_matrix = {}  # 创建初始邻接矩阵
    for i in matrix_items:
        adjacency_matrix[i] = {}
        for j in matrix_items:
            adjacency_matrix[i][j] = 0
    for i in adjacency_list:
        adjacency_matrix[i[0]][i[1]] = 1
        adjacency_matrix[i[1]][i[0]] = 1
    adjacency_matrix = pd.DataFrame(adjacency_matrix)
    index = adjacency_matrix.index.tolist()
    columns = adjacency_matrix.columns.tolist()
    np_adjacency_matrix = np.asarray(adjacency_matrix)
    adjacencies = {}
    for i in range(len(index)):
        for j in range(len(index)):
            if int(np_adjacency_matrix[i][j]) == 1:
                if '$' == index[i] or '$' == index[j]:
                    continue
                pair = sorted([index[i], index[j]])
                key = pair[0] + '@' + pair[1]
                if key not in adjacencies.keys():
                    adjacencies[key] = 1
                else:
                    adjacencies[key] += 1
    adjs = {}
    for i in adjacencies.keys():
        itemset = i.split('@')
        if itemset[0] not in adjs.keys():
            adjs[itemset[0]] = itemset[1]
        if itemset[1] not in adjs.keys():
            adjs[itemset[1]] = itemset[0]
    startpoint = []
    # 遍历矩阵第一行，不是0的加入start:
    for j in range(len(columns)):
        if np_adjacency_matrix[0][j] == 1:
            startpoint.append(columns[j])
    markerstartpoint = []
    chr = []
    for i in startpoint:
        if i not in markerstartpoint:
            path = []
            if i[-1] == 'a':
                path.append(i[:-1])
            else:
                path.append('-' + i[:-1])
            start = reverse(i)
            if start in startpoint:
                markerstartpoint.append(start)
                chr.append(path)
            else:
                while True:
                    next = adjs[start]
                    if next[-1] == 'a':
                       path.append(next[:-1])
                    else:
                        path.append('-' + next[:-1])
                    start = reverse(next)
                    if start in startpoint:
                        markerstartpoint.append(start)
                        break
                chr.append(path)
    vector = []
    for i in chr:
        for j in i:
            if j.startswith('-'):
                vector.append(j[1:])
            else:
                vector.append(j)
    cyclepoint = []
    for i in adjs.keys():
        if i[:-1] not in vector:
            cyclepoint.append(i)
    cyclechr = []
    markercycle = []
    for i in cyclepoint:
        if i not in markercycle:
            startpoint = i
            cycle = []
            markercycle.append(i)
            start = i
            while True:
                next = adjs[start]
                if next[-1] == 'a':
                    cycle.append(next[:-1])
                else:
                    cycle.append('-' + next[:-1])
                markercycle.append(start)
                markercycle.append(next)
                start = reverse(next)
                if start == startpoint:
                    break
            cyclechr.append(cycle)
    return chr,cyclechr

def changeAdj(adj_list):
    change = copy.deepcopy(adj_list)
    endpoints = []
    for j in change:
        for k in j:
            endpoints.append(k)
    random.shuffle(endpoints)
    change_part = []
    for j in range(int(len(endpoints) / 2)):
        change_part.append([endpoints[j * 2], endpoints[2 * j + 1]])
    return change_part

def buildNoCRBSimulations(prefix_adjacencies,
                          suffix_adjacencies, save_final_species_adjacencies, change_adjacency_number,
                          divergence_level, current_level):
    split_adjacency = []
    divergence_change_adjacencies_group_number = len(suffix_adjacencies) * 2
    for i in suffix_adjacencies:
        one_change_adjacencies = []
        for j in range(divergence_change_adjacencies_group_number):
            one_change_adjacencies.append(copy.deepcopy(i[j * change_adjacency_number:(j + 1) * change_adjacency_number]))
        one_change_adjacencies.append(copy.deepcopy(i[divergence_change_adjacencies_group_number * change_adjacency_number:]))
        split_adjacency.append(one_change_adjacencies)

    sub_species_1 = []
    sub_species_2 = []
    copy_number = 0

    for i in split_adjacency:
        sp1_copy = []
        sp2_copy = []
        sp1_change = copy_number
        sp2_change = copy_number+len(split_adjacency)
        for j in range(len(i)):
            # change different part, no CRBs
            if j == sp1_change:
                change = copy.deepcopy(i[j])
                change_part = changeAdj(change)
                sp1_copy.append(change_part)
            else:
                sp1_copy.append(copy.deepcopy(i[j]))
            if j == sp2_change:
                change = copy.deepcopy(i[j])
                change_part = changeAdj(change)
                sp2_copy.append(change_part)
            else:
                sp2_copy.append(copy.deepcopy(i[j]))

        sub_species_1.append(sp1_copy)
        sub_species_2.append(sp2_copy)
        copy_number += 1

    species_1 = []
    for i in range(len(sub_species_1)):
        one_copy = []
        one_copy += copy.deepcopy(prefix_adjacencies[i])
        for j in copy.deepcopy(sub_species_1[i]):
            one_copy += j
        species_1.append(one_copy)
    # save first species
    save_final_species_adjacencies.append(species_1)
    species_2 = []
    for i in range(len(sub_species_2)):
        one_copy = []
        one_copy += copy.deepcopy(prefix_adjacencies[i])
        for j in copy.deepcopy(sub_species_2[i]):
            one_copy += j
        species_2.append(one_copy)
    # save second species
    save_final_species_adjacencies.append(species_2)

    # duplication
    dup_flag = 1
    species_2_dup = []
    if dup_flag == 1:
        for i in range(len(sub_species_2)):
            change_list = copy.deepcopy(sub_species_2[i][:-1])
            unchange = copy.deepcopy(sub_species_2[i][-1])

            split_unchange = []
            for j in range(len(sub_species_2)*2):
                split_unchange.append(copy.deepcopy(unchange[j * change_adjacency_number:(j + 1) * change_adjacency_number]))
            split_unchange.append(unchange[len(sub_species_2) * 2 * change_adjacency_number:])
            # change adjacencies
            change_1 = i*len(sub_species_2)
            change_2 = i*len(sub_species_2) + 1
            new_change_list_copy1 = []
            new_change_list_copy2 = []
            for j in range(len(split_unchange)):
                if j == change_1:
                    change = copy.deepcopy(split_unchange[j])
                    change_part = changeAdj(change)
                    new_change_list_copy1.append(change_part)
                else:
                    new_change_list_copy1.append(copy.deepcopy(split_unchange[j]))
                if j == change_2:
                    change = copy.deepcopy(split_unchange[j])
                    change_part = changeAdj(change)
                    new_change_list_copy2.append(change_part)
                else:
                    new_change_list_copy2.append(copy.deepcopy(split_unchange[j]))
            final_change_list_1 = copy.deepcopy(change_list) + new_change_list_copy1
            final_change_list_2 = copy.deepcopy(change_list) + new_change_list_copy2
            species_2_dup.append(final_change_list_1)
            species_2_dup.append(final_change_list_2)

    if current_level == divergence_level:
        species_2_dup_sequence = []
        for i in range(len(species_2_dup)):
            one_copy = []
            one_copy += copy.deepcopy(prefix_adjacencies[int(i / 2)])
            for j in copy.deepcopy(species_2_dup[i]):
                one_copy += j
            species_2_dup_sequence.append(one_copy)
        save_final_species_adjacencies.append(species_2_dup_sequence)

    else:
        species_2_dup_sequence = []
        for i in range(len(species_2_dup)):
            one_copy = []
            one_copy += copy.deepcopy(prefix_adjacencies[int(i / 2)])
            for j in copy.deepcopy(species_2_dup[i]):
                one_copy += j
            species_2_dup_sequence.append(one_copy)
        # save duplicated species
        save_final_species_adjacencies.append(species_2_dup_sequence)
        new_prefix_adjacency = []
        new_suffix_adjacency = []
        for i in range(len(species_2_dup)):
            one_copy = copy.deepcopy(prefix_adjacencies[int(i / 2)])
            change_part = species_2_dup[i][:-1]
            unchange_part = species_2_dup[i][-1]
            for j in change_part:
                one_copy += copy.deepcopy(j)
            new_prefix_adjacency.append(one_copy)
            new_suffix_adjacency.append(copy.deepcopy(unchange_part))
        # recursion build next level species, level += 1
        buildNoCRBSimulations(new_prefix_adjacency, new_suffix_adjacency,
                              save_final_species_adjacencies, change_adjacency_number,
                              divergence_level, current_level + 1)

def simulateNoCRB(workdir):
    chromosome_number = 5
    block_number = 100
    ancestor_sequence = []
    one_chromosome = int(block_number / chromosome_number)
    block = 100
    for i in range(chromosome_number):
        sequence = []
        for j in range(one_chromosome):
            if block % 2 == 0:
                sequence.append('-' + str(block))
            else:
                sequence.append(str(block))
            block += 1
        ancestor_sequence.append(sequence)
    ancestor_adjacency = sequence2adjacency(ancestor_sequence)
    random.shuffle(ancestor_adjacency)
    print('ancestor adjacency number:')
    print(len(ancestor_adjacency))
    divergence_level = 2
    change_adjacency_number = 5
    save_final_species_adjacencies = []
    buildNoCRBSimulations([[]], [ancestor_adjacency],
                          save_final_species_adjacencies, change_adjacency_number, divergence_level, current_level=0)
    species_count = 1
    # output species
    for i in save_final_species_adjacencies:
        copy_count = 1
        outfile = workdir + 'species.sequence.' + str(species_count)

        outfile = open(outfile,'w')
        for j in i:
            filter_tel2tel = []
            for k in j:
                # filter ($,$)
                if k[0] == '$' and k[1] == '$':
                    continue
                else:
                    if k[0] != '$':
                        newendpoint1 = k[0][:-1]+'_'+str(copy_count)+k[0][-1]
                    else:
                        newendpoint1 = '$'
                    if k[1] != '$':
                        newendpoint2 = k[1][:-1]+'_'+str(copy_count)+k[1][-1]
                    else:
                        newendpoint2 = '$'
                    filter_tel2tel.append([newendpoint1,newendpoint2])
            chrs,cycles = assemble(filter_tel2tel)
            for k in chrs:
                outfile.write('s ')
                for l in k:
                    outfile.write(l+' ')
                outfile.write('\n')
            for k in cycles:
                outfile.write('c ')
                min_index = -1
                min_value = 1000000
                for l in range(len(k)):
                    if k[l].startswith('-'):
                        item = k[l][1:].split('_')
                        block = int(item[0])
                    else:
                        item = k[l].split('_')
                        block = int(item[0])
                    if block < min_value:
                        min_index = l
                        min_value = block
                if k[min_index].startswith('-'):
                    half1 = k[min_index + 1:]
                    half2 = k[:min_index + 1]
                    new_string = half1 + half2
                else:
                    half1 = k[min_index:]
                    half2 = k[:min_index]
                    new_string = half1 + half2
                for l in new_string:
                    outfile.write(l+' ')
                outfile.write('\n')
            copy_count += 1
        outfile.close()
        species_count += 1

def doubled(infile,outfile):
    outfile = open(outfile,'w')
    sequence = []
    with open(infile,'r') as f:
        while True:
            line = f.readline()
            if not line:
                break
            sequence.append(line)
    for i in sequence:
        outfile.write(i)
    for i in sequence:
        outfile.write(i)

workdir = 'D:/InferAncestorGenome/realData/IAGS_version1.0/simulations/NonCRBs/inferring/repeat/'
# simulate No CRB species
if not os.path.exists(workdir):
    os.makedirs(workdir)
resultfile = open(workdir + 'result.xls', 'w')
resultfile.write('Repeat\tAncestor\tFissions\tFusions\n')
for i in range(200):
    repeatdir = workdir + str(i+1) + '/'
    if not os.path.exists(repeatdir):
        os.makedirs(repeatdir)
    simulateNoCRB(repeatdir)
    filelist = ['species.sequence.1', 'species.sequence.2', 'species.sequence.3',
                'species.sequence.4', 'species.sequence.5', 'species.sequence.6',
                'species.sequence.7', 'species.sequence.8', 'species.sequence.9']
    for j in filelist:
        sequence = readSequence(repeatdir + j)
        outSequence(sequence, repeatdir + j + '.noBar')
    """
    Ancestor 8: Multi-copy GGHP model
    """
    dup_child_file = repeatdir + 'species.sequence.9.noBar'
    outgroup_file = repeatdir + 'species.sequence.7.noBar'
    outAncestor8dir = repeatdir + 'Ancestor8/'
    if not os.path.exists(outAncestor8dir):
        os.makedirs(outAncestor8dir)

    dup_copy_number = 8
    out_copy_number = 4
    ancestor_target_copy_number = 4
    ancestor_name = 'Ancestor8'
    MultiCopyGGHPmodel(dup_child_file, outgroup_file, outAncestor8dir,
                       ancestor_name, dup_copy_number, out_copy_number,
                       ancestor_target_copy_number)

    ancestor_file = outAncestor8dir + ancestor_name + '.block'

    fissions, fusions = calculateFissionAndFussions(ancestor_file,
                                                    repeatdir + 'species.sequence.8.noBar',
                                                    ancestor_target_copy_number, ancestor_target_copy_number,
                                                    outAncestor8dir)
    resultfile.write(str(i+1)+'\tAncestor8\t' + str(fissions) + '\t' + str(fusions) + '\n')

    """
    Ancestor 5: Multi-copy GGHP model
    """
    dup_child_file = repeatdir + 'species.sequence.7.noBar'
    outgroup_file = repeatdir + 'species.sequence.4.noBar'
    outAncestor5dir = repeatdir + 'Ancestor5/'
    if not os.path.exists(outAncestor5dir):
        os.makedirs(outAncestor5dir)
    dup_copy_number = 4
    out_copy_number = 2
    ancestor_target_copy_number = 2
    ancestor_name = 'Ancestor5'
    MultiCopyGGHPmodel(dup_child_file, outgroup_file, outAncestor5dir,
                       ancestor_name, dup_copy_number, out_copy_number,
                       ancestor_target_copy_number)

    ancestor_file = outAncestor5dir + ancestor_name + '.block'

    fissions, fusions = calculateFissionAndFussions(ancestor_file,
                                                    repeatdir + 'species.sequence.5.noBar',
                                                    ancestor_target_copy_number, ancestor_target_copy_number,
                                                    outAncestor5dir)
    resultfile.write(str(i+1)+'\tAncestor5\t' + str(fissions) + '\t' + str(fusions) + '\n')

    """
    Ancestor 6: Multi-copy GMP model
    """
    outAncestor6dir = repeatdir + 'Ancestor6/'
    if not os.path.exists(outAncestor6dir):
        os.makedirs(outAncestor6dir)
    doubled(outAncestor5dir + 'Ancestor5.block', outAncestor6dir + 'Ancestor5.doubled.block')
    species_file_list = [repeatdir + 'species.sequence.7.noBar',
                         outAncestor8dir + 'Ancestor8.block',
                         outAncestor6dir + 'Ancestor5.doubled.block']
    guided_species_for_matching = repeatdir + 'species.sequence.7.noBar'
    ancestor_target_copy_number = 4
    ancestor_name = 'Ancestor6'
    MultiCopyGMPmodel(species_file_list, outAncestor6dir, guided_species_for_matching,
                      ancestor_name, ancestor_target_copy_number)

    ancestor_file = outAncestor6dir + ancestor_name + '.block'

    fissions, fusions = calculateFissionAndFussions(ancestor_file,
                                                    repeatdir + 'species.sequence.6.noBar',
                                                    ancestor_target_copy_number, ancestor_target_copy_number,
                                                    outAncestor6dir)
    resultfile.write(str(i+1)+'\tAncestor6\t' + str(fissions) + '\t' + str(fusions) + '\n')

    """
    Ancestor 2: GGHP model
    """
    dup_child_file = repeatdir + 'species.sequence.4.noBar'
    outgroup_file = repeatdir + 'species.sequence.1.noBar'
    outAncestor2dir = repeatdir + 'Ancestor2/'
    if not os.path.exists(outAncestor2dir):
        os.makedirs(outAncestor2dir)
    dup_copy_number = 2
    out_copy_number = 1
    ancestor_target_copy_number = 1
    ancestor_name = 'Ancestor2'
    GGHPmodel(dup_child_file=dup_child_file,
              outgroup_file=outgroup_file,
              outdir=outAncestor2dir,
              ancestor_name=ancestor_name,
              dup_copy_number=dup_copy_number,
              out_copy_number=out_copy_number)

    ancestor_file = outAncestor2dir + ancestor_name + '.block'

    fissions, fusions = calculateFissionAndFussions(ancestor_file,
                                                    repeatdir + 'species.sequence.2.noBar',
                                                    ancestor_target_copy_number, ancestor_target_copy_number,
                                                    outAncestor2dir)
    resultfile.write(str(i+1)+'\tAncestor2\t' + str(fissions) + '\t' + str(fusions) + '\n')

    """
    Ancestor 3: Multi-copy GMP model
    """
    outAncestor3dir = repeatdir + 'Ancestor3/'
    outAncestor2dir = repeatdir + 'Ancestor2/'
    if not os.path.exists(outAncestor3dir):
        os.makedirs(outAncestor3dir)
    doubled(outAncestor2dir + 'Ancestor2.block', outAncestor3dir + 'Ancestor2.doubled.block')
    species_file_list = [repeatdir + 'species.sequence.4.noBar',
                         outAncestor5dir + 'Ancestor5.block',
                         outAncestor3dir + 'Ancestor2.doubled.block']
    guided_species_for_matching = repeatdir + 'species.sequence.4.noBar'
    ancestor_target_copy_number = 2
    ancestor_name = 'Ancestor3'
    MultiCopyGMPmodel(species_file_list, outAncestor3dir, guided_species_for_matching,
                      ancestor_name, ancestor_target_copy_number)

    ancestor_file = outAncestor3dir + ancestor_name + '.block'

    fissions, fusions = calculateFissionAndFussions(ancestor_file,
                                                    repeatdir + 'species.sequence.3.noBar',
                                                    ancestor_target_copy_number, ancestor_target_copy_number,
                                                    outAncestor3dir)
    resultfile.write(str(i+1)+'\tAncestor3\t' + str(fissions) + '\t' + str(fusions) + '\n')
    resultfile.flush()

resultfile.close()


#!/usr/bin/env python
# TODO: how do we choose the simulation box?

import numpy as np
import numpy.linalg as la
import sys
import os
from libs import topology as top
from libs import base
from libs import utils

class Options(object):

	def __init__(self):
		object.__init__(self)

		self.closed = True
		self.double = True
		self.nicked = False
		self.supercoiling = 0.
		self.writhe = 0.
		self.seed = None
		self.sequence_file = None

	def check(self):
		if self.nicked and not self.double:
			print >> sys.stderr, "The --nicked and --ssDNA options are incompatible"
			exit(1)

def print_usage():
	print >> sys.stderr, "USAGE:"
	print >> sys.stderr, "\t%s centerline_file" % sys.argv[0]
	print >> sys.stderr, "\t[-c\--closed] [-o\--open] [-h\--help] [-d\--dsDNA] [-s\--ssDNA] [-n\--nicked] [-p\--supercoiling VALUE] [-w\--writhe VALUE] [-e\--seed VALUE] [-q\--sequence FILE]"
	exit(1)

def parse_options():
	shortArgs = 'cohdsnp:w:e:q:'
	longArgs = ['closed', 'open', 'help', 'dsDNA', 'ssDNA', 'nicked', 'supercoiling=', 'writhe=', 'seed=', 'sequence=']

	opts = Options()

	try:
		import getopt
		args, files = getopt.gnu_getopt(sys.argv[1:], shortArgs, longArgs)
		for k in args:
			if k[0] == '-c' or k[0] == '--closed': opts.closed = True
			if k[0] == '-o' or k[0] == '--open': opts.closed = False
			if k[0] == '-h' or k[0] == '--help': print_usage()
			if k[0] == '-d' or k[0] == '--dsDNA': opts.double = True
			if k[0] == '-s' or k[0] == "--ssDNA": opts.double = False
			if k[0] == '-n' or k[0] == "--nicked": opts.nicked = True
			if k[0] == '-p' or k[0] == "--supercoiling": opts.supercoiling = float(k[1])
			if k[0] == '-w' or k[0] == "--writhe": opts.writhe = float(k[1])
			if k[0] == '-e' or k[0] == "--seed": opts.seed = int(k[1])
			if k[0] == '-q' or k[0] == "--sequence": opts.sequence_file = k[1]

		opts.centerline_file = files[0]
	except Exception:
		print_usage()

	return opts

# base-base distance along the helical pitch
BASE_BASE = 0.3897628551303122
# distance between the helix centre and the nucleotides' centre of mass
CM_CENTER_DS = 0.6

if __name__ == '__main__':
	opts = parse_options()
	opts.check()

	if opts.seed != None:
		np.random.seed(opts.seed)

	# import the coordinates from the user-provided file
	coordxyz = np.loadtxt(opts.centerline_file, float)

	# number of base pairs
	nbases = len(coordxyz)

	# use the model parameters to scale the distances
	scaling = BASE_BASE / la.norm(coordxyz[1, :] - coordxyz[0, :])
	coordxyz *= scaling

	#initialize vectors
	dist = np.copy(coordxyz)
	dist_norm = np.copy(coordxyz)
	p = np.copy(coordxyz)

	ssdna1 = np.copy(coordxyz)
	v_perp_ssdna1 = np.copy(coordxyz)

	ssdna2 = np.copy(coordxyz)
	v_perp_ssdna2 = np.copy(coordxyz)

	# take the bounding box as the simulation box for the output configuration
	boxx = max(coordxyz[:nbases, 0]) - min(coordxyz[:nbases, 0])
	boxy = max(coordxyz[:nbases, 1]) - min(coordxyz[:nbases, 1])
	boxz = max(coordxyz[:nbases, 2]) - min(coordxyz[:nbases, 2])
	boxmax = 1.5 * max(boxx, boxy, boxz)

	# centerline base_to_base vectors
	for c in range(0, nbases):
		ind = c
		ind1 = (c + 1) % nbases
		#open chain cannot compute dist at c=nbases-1
		#so we use as reference the previus dist
		if opts.closed or (not opts.closed and c!=nbases-1):
			dist[ind, :] = coordxyz[ind1, :] - coordxyz[ind, :]
			dist_norm[ind, :] = dist[ind, :] / la.norm(dist[ind, :])
		else:
			dist[ind, :] = dist[ind-1, :]
			dist_norm[ind, :] = dist_norm[ind-1, :]

	# vectors perpendicular between two consecutive centerline vectors (normalized)
	for c in range(0, nbases):
		ind_1 = (c - 1 + nbases) % nbases
		ind = c

		#opens chain have random p at c=0 and at c=nbases-1 due to absence of neighbours
		if opts.closed or (not opts.closed and c!=0 and c!=nbases-1):
			#check that dist[ind_1, :] and dist[ind, :] are not equals
			if not np.all(np.isclose( dist[ind_1, :] , dist[ind, :]) ):
				p[ind, :] = np.cross(dist[ind_1, :] , dist[ind, :])
				p[ind, :] /= la.norm(p[ind, :])
			#else assign random p
			else:
				rv=np.random.uniform(-1,1,3)
				p[ind, :] = rv - (np.dot(dist_norm[ind, :],rv)) * dist_norm[ind, :]
				p[ind, :] /= la.norm(p[ind, :])
		else:
			rv=np.random.uniform(-1,1,3)
			p[ind, :] = rv - (np.dot(dist_norm[ind, :],rv)) * dist_norm[ind, :]
			p[ind, :] /= la.norm(p[ind, :])

	# chain writhe
	WR = 0.
	if opts.closed:
		WR = top.get_writhe(coordxyz)

	#oxdna equiibrium pitch
	pitch = 10.5

	#global linking number
	LK = round((nbases / pitch) * (opts.supercoiling + 1) + opts.writhe)

	TW = LK - WR

	#twisting angle between two consecutive bases
	rot_base = TW * 2.0 * np.pi / nbases

	#recap
	if opts.closed:
		print >> sys.stderr, "Total Linking Number (LK) %f, composed of:" % ((nbases / pitch) * (opts.supercoiling + 1) + opts.writhe)
		print  >> sys.stderr, "1) Equilibrium number of DNA turns %f" % (nbases / pitch)
		print  >> sys.stderr, "2) Target writhe %f = Topological writhe %f + Turns imposed by supercoiling  %f " % (opts.writhe+(nbases / pitch) *opts.supercoiling,opts.writhe,(nbases / pitch) *opts.supercoiling)
		print >> sys.stderr, "LK has been rounded to %f" % LK
		print >> sys.stderr, "Initial chain writhe %f" % WR

	####################################
	# Initialize DNA strand
	####################################

	#First hydrogen-hydrogen bond vector
	v_perp_ssdna1[0, :] = np.cross(dist_norm[0, :], p[0, :])
	v_perp_ssdna1[0, :] /= la.norm(v_perp_ssdna1[0, :])
	v_perp_ssdna2[0, :] = -v_perp_ssdna1[0, :]

	for c in range(nbases):
		#dna center of mass positions
		ssdna1[c, :] = coordxyz[c, :] - CM_CENTER_DS * v_perp_ssdna1[c, :]
		ssdna2[c, :] = coordxyz[c, :] + CM_CENTER_DS * v_perp_ssdna1[c, :]

		#Update v_perp_ssdna1

		ind_1 = c
		ind = (c + 1) % nbases

		alpha = top.py_ang(v_perp_ssdna1[ind_1, :] , p[ind, :] , dist[ind_1, :])
		gamma = rot_base - alpha
		#prevent change when ind=0 on open chain
		if c!=nbases-1:
			R = utils.get_rotation_matrix(dist[ind, :], gamma)
			v_perp_ssdna1[ind, :] = np.dot(R , p[ind, :])
			v_perp_ssdna2[ind, :] = -v_perp_ssdna1[ind, :]

	# check LK imposed and measured
	if opts.closed:
		TW_measured = top.get_twist(coordxyz, ssdna1)
	box = np.array([boxmax, boxmax, boxmax])
	system = base.System(box)

	if opts.sequence_file == None:
		ssdna1_base = np.zeros(nbases, int)
		for c in range(nbases):
			ssdna1_base[c] = np.random.randint(0, 4)
	else:
		try:
			seq_file = open(opts.sequence_file)
		except Exception:
			print >> sys.stderr, "The sequence file '%s' is unreadable" % opts.sequence_file
			exit(1)

		contents = seq_file.read()
		# remove all whitespace from the file's contents
		sequence = ''.join(contents.split())
		if len(sequence) != nbases:
			print >> sys.stderr, "The length of the given sequence (%d) should be equal to the number of coordinates in the centerline file (%d)" % (len(sequence), nbases)
			exit(1)

		ssdna1_base = map(lambda x: base.base_to_number[x], sequence)

		seq_file.close()

	strand1 = base.Strand()
	for c in range(nbases):
		b = ssdna1_base[c]
		strand1.add_nucleotide(base.Nucleotide(ssdna1[c], v_perp_ssdna1[c], dist_norm[c], b, b))
	if opts.closed:
		strand1.make_circular()
	system.add_strand(strand1)

	if opts.double:
		strand2 = base.Strand()
		for c in range(nbases):
			reverse_idx = nbases - 1 - c
			b = 3 - ssdna1_base[reverse_idx]
			strand2.add_nucleotide(base.Nucleotide(ssdna2[reverse_idx], v_perp_ssdna2[reverse_idx], -dist_norm[reverse_idx], b, b))
		if opts.closed and not opts.nicked:
			strand2.make_circular()
		system.add_strand(strand2)

	basename = os.path.basename(sys.argv[1])
	topology_file = basename + ".top"
	configuration_file = basename + ".oxdna"
	system.print_lorenzo_output(configuration_file, topology_file)

	print >> sys.stderr, "## Wrote data to '%s' / '%s'" % (configuration_file, topology_file)
	print >> sys.stderr, "## DONE"


#-------------------------------------------------------------------------------
#
#  IGRF format parser
#
# Author: Martin Paces <martin.paces@eox.at>
#
#-------------------------------------------------------------------------------
# Copyright (C) 2018 EOX IT Services GmbH
#
# 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 of this Software or works derived from this 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.
#-------------------------------------------------------------------------------

from numpy import array

IGRF_EXTRAPOLATION_PERIOD = 5.0 # years

def parse_igrf_file(file_in):
    """ Parse IGRF file format and return a dictionary containing the parsed
    model data.
    """
    lines = strip_igrf_comments(file_in)
    data = {}
    data["labels"] = parse_igrf_header(next(lines))
    data["t"] = parse_igrf_times(next(lines))
    data["nm"], data["gh"] = parse_igrf_coefficients(lines)
    data["degree_min"] = data["nm"][:, 0].min()
    data["degree_max"] = data["nm"][:, 0].max()
    return data

def parse_igrf_coefficients(lines):
    """ Parse IGRF coefficients. """
    nm_idx = []
    coeff = []
    for line in lines:
        fields = line.split()
        label, n_idx, m_idx = fields[0], int(fields[1]), int(fields[2])
        if label == "h":
            m_idx = -m_idx
        nm_idx.append((n_idx, m_idx))
        coeff.append([float(v) for v in fields[3:]])
    coeff = array(coeff)
    coeff[:, -1] = coeff[:, -2] + coeff[:, -1] * IGRF_EXTRAPOLATION_PERIOD
    return array(nm_idx), coeff

def parse_igrf_times(line):
    """ Parse SHC times. """
    times = [float(v) for v in line.split()[3:-1]]
    times.append(times[-1] + IGRF_EXTRAPOLATION_PERIOD)
    return array(times)

def parse_igrf_header(line):
    """ Parse IGRF header with the column labels. """
    return line.split()

def strip_igrf_comments(file_in):
    """ Strip initial comments and empty lines from a text file stream. """
    for line in file_in:
        line = line.partition("#")[0].strip()
        if line:
            yield line
            break
    for line in file_in:
        line = line.strip()
        yield line


#! /usr/bin/env python
"""
agilent4155_matlab_param.py
Prarmeter extractor for matlab generated .xlsx files

Created by Jeremy Smith on 2015-07-07
University of California, Berkeley
j-smith@eecs.berkeley.edu
"""

import os
import sys
import xlrd
import numpy as np
import myfunctions as mf
from scipy import stats

__author__ = "Jeremy Smith"
__version__ = "1.5"

data_path = os.path.dirname(__file__)    # Path name for location of script

files = os.listdir(data_path)   # All files in directory
data_summary = []
colheads = ['VG', 'ID1', 'ID2', 'IG1', 'IG2']
skipinit = 3          # skip initial data points
rangefittop = 20.0    # % from max for fit range
rangefitbot = 20.0    # % from min for fit range
mobilitycorrection = 1.000     # K_actual/K_file
summary_list_header = [["filename", "satmob_tmax", "vthsat_tmax", "satmob_rev_tmax", "vthsat_rev_tmax",
						"linmob_tmax", "vthlin_tmax", "linmob_rev_tmax", "vthlin_rev_tmax",
						"hysteresis_lin", "hysteresis_sat",
						"onoffratio_lin", "onoffratio_sat",
						"leakage_ratio_lin", "leakage_ratio_sat",
						"Slin", "Slin_rev", "Ssat", "Ssat_rev",
						"satmob_FITTED", "vthsat_FITTED", "r_value"]]
sweepfwddirection = True       # True if sweeping from depletion to accumulation

def main():
	"""Main function"""

	print "\nBatch importing .xlsx files..."
	print data_path, '\n'

	for f in files:
		print f
		# Loops through all transfer files
		if "IDVG.xlsx" in f:

			workbook = xlrd.open_workbook(f, logfile=open(os.devnull, 'w'))

			for dev in workbook.sheet_names():
				if "Sheet" in dev:
					continue
				print "  - device {:s}".format(dev)
				datasheet = workbook.sheet_by_name(dev)
				run_numbers = [str(int(x)) for x in datasheet.row_values(2) if x]

				for i, run in enumerate(run_numbers):
					print "    - run {:s}".format(run)
					data = {}
					# File name for outputs
					outname = f[:-5] + '_' + dev + '_' + run
					# Constant parameters taken from header of .xlsx file
					vd1 = float(datasheet.cell_value(3, 6*i + 3))
					vd2 = float(datasheet.cell_value(4, 6*i + 3))
					chl = float(datasheet.cell_value(1, 1))
					chw = float(datasheet.cell_value(0, 1))
					tox = float(datasheet.cell_value(1, 3))
					kox = float(datasheet.cell_value(0, 3))
					ldr = float(datasheet.cell_value(1, 5))
					lso = float(datasheet.cell_value(0, 5))
					# Calculation of geometric capacitance
					ci = 8.85418782e-7*kox/tox
					# Extract data
					for h in colheads:
						data[h] = []
					for row in range(datasheet.nrows - 9):
						for col, h in enumerate(colheads):
							if datasheet.cell_type(9 + row, 6*i + col) is 0:
								continue
							data[h].append(float(datasheet.cell_value(9 + row, 6*i + col)))
					p = len(data['VG'])/2
					# Forward scan
					if sweepfwddirection:
						vg = np.array(data['VG'][:p])
						id1 = np.array(data['ID1'][:p])
						id2 = np.array(data['ID2'][:p])
						ig1 = np.array(data['IG1'][:p])
						ig2 = np.array(data['IG2'][:p])
					else:
						vg = np.array(data['VG'][:p][::-1])
						id1 = np.array(data['ID1'][:p][::-1])
						id2 = np.array(data['ID2'][:p][::-1])
						ig1 = np.array(data['IG1'][:p][::-1])
						ig2 = np.array(data['IG2'][:p][::-1])
					# Reverse scan
					if sweepfwddirection:
						vg_r = np.array(data['VG'][p:][::-1])
						id1_r = np.array(data['ID1'][p:][::-1])
						id2_r = np.array(data['ID2'][p:][::-1])
						ig1_r = np.array(data['IG1'][p:][::-1])
						ig2_r = np.array(data['IG2'][p:][::-1])
					else:
						vg_r = np.array(data['VG'][p:])
						id1_r = np.array(data['ID1'][p:])
						id2_r = np.array(data['ID2'][p:])
						ig1_r = np.array(data['IG1'][p:])
						ig2_r = np.array(data['IG2'][p:])

					# Smoothing Id for fitting
					id1_smoothed = mf.adjAvSmooth(abs(id1), N=1)
					id2_smoothed = mf.adjAvSmooth(abs(id2), N=1)
					id1_r_smoothed = mf.adjAvSmooth(abs(id1_r), N=1)
					id2_r_smoothed = mf.adjAvSmooth(abs(id2_r), N=1)
					# On-off ratio
					onoffratio1 = np.log10(max(id1[skipinit:-1])/min(abs(id1[skipinit:-1])))
					onoffratio2 = np.log10(max(id2[skipinit:-1])/min(abs(id2[skipinit:-1])))
					# Leakage ratio
					leakage_ratio1 = np.log10(abs(id1/ig1))
					leakage_ratio2 = np.log10(abs(id2/ig2))

					# Finding max saturation transconductance
					sqrtid2 = np.sqrt(id2_smoothed)
					sqrtid2_r = np.sqrt(id2_r_smoothed)
					diff_sqrt_id2_smoothed = np.array(mf.numDiff(sqrtid2, vg))
					diff_sqrt_id2_r_smoothed = np.array(mf.numDiff(sqrtid2_r, vg_r))
					tsmaxarg = np.argmax(diff_sqrt_id2_smoothed[skipinit:-1]) + skipinit
					tsmaxarg_r = np.argmax(diff_sqrt_id2_r_smoothed[skipinit:-1]) + skipinit
					# Saturation mobility (max transconductance)
					satmob_t = mobilitycorrection * (2*chl/(chw*ci))*(diff_sqrt_id2_smoothed)**2
					satmob_r_t = mobilitycorrection * (2*chl/(chw*ci))*(diff_sqrt_id2_r_smoothed)**2
					satmob_tmax = satmob_t[tsmaxarg]
					satmob_r_tmax = satmob_r_t[tsmaxarg_r]
					# Saturation threshold voltage (max transconductance)
					vthsat_tmax = vg[tsmaxarg] - sqrtid2[tsmaxarg]/diff_sqrt_id2_smoothed[tsmaxarg]
					vthsat_r_tmax = vg_r[tsmaxarg_r] - sqrtid2_r[tsmaxarg_r]/diff_sqrt_id2_r_smoothed[tsmaxarg_r]
					# Hysteresis
					hysteresissat = vthsat_tmax - vthsat_r_tmax
					# Calculate subthreshold slopes
					sts_sat = min(abs(1/np.array(mf.numDiff([np.log10(abs(x)) for x in id2_smoothed[skipinit:-1]], vg[skipinit:-1]))))
					sts_r_sat = min(abs(1/np.array(mf.numDiff([np.log10(abs(x)) for x in id2_r_smoothed[skipinit:-1]], vg_r[skipinit:-1]))))

					# Finding max linear transconductance
					diff_id1_smoothed = np.array(mf.numDiff(id1_smoothed, vg))
					diff_id1_r_smoothed = np.array(mf.numDiff(id1_r_smoothed, vg_r))
					tlmaxarg = np.argmax(diff_id1_smoothed[skipinit:-1]) + skipinit
					tlmaxarg_r = np.argmax(diff_id1_r_smoothed[skipinit:-1]) + skipinit
					# Linear mobility (max transconductance)
					linmob_t = mobilitycorrection * (chl/(chw*ci*vd1))*(diff_id1_smoothed)
					linmob_r_t = mobilitycorrection * (chl/(chw*ci*vd1))*(diff_id1_r_smoothed)
					linmob_tmax = linmob_t[tlmaxarg]
					linmob_r_tmax = linmob_r_t[tlmaxarg_r]
					# Linear threshold voltage (max transconductance)
					vthlin_tmax = vg[tlmaxarg] - id1_smoothed[tlmaxarg]/diff_id1_smoothed[tlmaxarg]
					vthlin_r_tmax = vg_r[tlmaxarg_r] - id1_r_smoothed[tlmaxarg_r]/diff_id1_r_smoothed[tlmaxarg_r]
					# Hysteresis
					hysteresislin = vthlin_tmax - vthlin_r_tmax
					# Calculate subthreshold slopes
					sts_lin = min(abs(1/np.array(mf.numDiff([np.log10(abs(x)) for x in id1_smoothed[skipinit:-1]], vg[skipinit:-1]))))
					sts_r_lin = min(abs(1/np.array(mf.numDiff([np.log10(abs(x)) for x in id1_r_smoothed[skipinit:-1]], vg_r[skipinit:-1]))))

					# Finds range of data that lies within the minimum+x% and the maximum-x% and also has a positive transconductance
					fitrange_id_lo = (1-rangefitbot/100.0)*min(sqrtid2[skipinit:-1]) + (rangefitbot/100.0)*max(sqrtid2[skipinit:-1])
					fitrange_id_hi = (1-rangefittop/100.0)*max(sqrtid2[skipinit:-1]) + (rangefittop/100.0)*min(sqrtid2[skipinit:-1])
					fitrange_bool = np.bitwise_and(np.bitwise_and(sqrtid2 > fitrange_id_lo, sqrtid2 < fitrange_id_hi), diff_sqrt_id2_smoothed > 0)
					# Checks that there are at least 3 data points to fit
					if sum(fitrange_bool) < 3:
						print "      NOT ENOUGH DATA TO FIT"
						satmob_FITTED = np.nan
						vthsat_FITTED = np.nan
						r_value = np.nan
					else:
						# Linear Fitting to sqrt(Idrain)
						slope, intercept, r_value, p_value, std_err = stats.linregress(vg[fitrange_bool][skipinit:-1], sqrtid2[fitrange_bool][skipinit:-1])
						fitline = slope*vg + intercept
						# Saturation mobility (from slope of sqrt(Idrain) fit)
						satmob_FITTED = mobilitycorrection * (2*chl/(chw*ci))*slope**2
						# Threshold Voltage (from slope of sqrt(Idrain) fit)
						vthsat_FITTED = -intercept/slope
						# Plot sqrt(Isd)
						mf.quickPlot(outname+"_SQRTplot", data_path, [vg, sqrtid2, fitline],
						xlabel="VG [V]", ylabel="sqrt(Id) [A^0.5]", yrange=[0, 'auto'])

					# Output data
					data_summary.append([outname,
						satmob_tmax, vthsat_tmax,
						satmob_r_tmax, vthsat_r_tmax,
						linmob_tmax, vthlin_tmax,
						linmob_r_tmax, vthlin_r_tmax,
						hysteresislin, hysteresissat,
						onoffratio1, onoffratio2,
						leakage_ratio1[-skipinit],
						leakage_ratio2[-skipinit],
						sts_lin, sts_r_lin, sts_sat, sts_r_sat,
						satmob_FITTED, vthsat_FITTED, r_value**2])

					# Ouput files
					mf.dataOutputHead(outname+"_transfer.txt", data_path, [np.array(data['VG']), abs(np.array(data['ID1'])), abs(np.array(data['ID2'])), abs(np.array(data['IG1'])), abs(np.array(data['IG2'])),
						np.concatenate((linmob_t, linmob_r_t[::-1])), np.concatenate((satmob_t, satmob_r_t[::-1]))],
						[["vg", "idlin", "idsat", "iglin", "igsat", "LINMOB", "SATMOB"]],
						format_d="%.3f\t %.5e\t %.5e\t %.5e\t %.5e\t %.5e\t %.5e\n",
						format_h="%s\t")

					# Plot transfer
					mf.quickPlot(outname+"_TRANSFERplot", data_path, [vg, id1_smoothed, id1_r_smoothed, abs(ig1), id2_smoothed, id2_r_smoothed, abs(ig2)],
						xlabel="VG [V]", ylabel="Id,g [A]", yscale="log", yrange=[1e-12, 1e-2], col=["r", "r", "r", "b", "b", "b"])

	mf.dataOutputHead("SUMMARY.txt", data_path, map(list, zip(*data_summary)), summary_list_header,
		format_d="%s\t %.5e\t %.5f\t %.5e\t %.5f\t %.5e\t %.5f\t %.5e\t %.5f\t %.5f\t %.5f\t %.5f\t %.5f\t %.5f\t %.5f\t %.5f\t %.5f\t %.5f\t %.5f\t %.5e\t %.5f\t %.6f\n",
		format_h="%s\t")

	return

if __name__ == "__main__":
	sys.exit(main())


# coding=utf-8
# Copyright (c) 2015-2018, UT-BATTELLE, LLC
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors
# may be used to endorse or promote products derived from this software without
# specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
CISM_glissade module for numerics analysis
"""
from __future__ import absolute_import, division, print_function, unicode_literals

import math

import numpy as np

from netCDF4 import Dataset
from scipy import interpolate

class DataGrid:
    """
    Class to handle the CISM_glissade grids, which are cell-centered grids.
    """
    def __init__(self, data):
        self.y = data.variables['y1']
        self.ny = self.y[:].shape[0]
        self.dy = self.y[1] - self.y[0]

        # NOTE: cell centered grids, hence the dy.
        self.Ly = self.y[-1] - self.y[0] + self.dy

        self.x = data.variables['x1']
        self.nx = self.x[:].shape[0]
        self.dx = self.x[1] - self.x[0]

        # NOTE: cell centered grids, hence the dx.
        self.Lx = self.x[-1] - self.x[0] + self.dx

        self.y_hat = (self.y[:] + self.y[0])/self.Ly
        self.x_hat = (self.x[:] + self.x[0])/self.Lx

class RotatedGrid:
    """
    For the ISMIP-HOM f tests, CISM computes the flow of a glacier down an
    inclined plane:

    ::

        z
        ^
        |
        * .
        .    .
        .        .
        *           .
        | .             .
        |    .     ICE      .
        |        .              .
        |            .             .
        |  BED           .            .
        |                    .           *
        |                       .        .
        |                          .     .
        |                             .  .
        |                                *
        |                                |
        |                                |
        |                                |
        0------------------------------------->x

    The origin is at point 0 in the above figure, and the topmost point of the
    glacier is at x=0, z=7 (in km). The slope is 3 degrees. The ice is 1000 m
    tall and flows down the inclined plane.

    ISMIP-HOM, however, defines the coordinate system with the origin located at
    the topmost point of the glacier (0,7) with the x' axis pointing down slope
    and z' pointing perpendicular to the slope. So the coordinate system is
    shifted, and rotated by a=3 degrees from the CISM glissade grid.

    An additional complication is that the surface is computed in CISM on the
    standard grid, but velocities are computed on a staggered, grid.

    This class converts the CISM_glissade coordinate system to the ISMIP-HOM
    coordinate system.
    """
    def __init__(self, alpha, data):
        self.alpha = alpha
        self.y0 = data.variables['y0'][:]
        self.x0 = data.variables['x0'][:]

        self.usurf_ustag = data.variables['usurf'][-1,:,:]
        self.usurf_stag = (  self.usurf_ustag[1: , 1: ] + self.usurf_ustag[1: , :-1]
                           + self.usurf_ustag[:-1, :-1] + self.usurf_ustag[:-1, 1: ]) / 4.0

        self.usurf = -(self.x0)*math.sin(alpha) + (self.usurf_stag-7000.0)*math.cos(alpha)

        self.uvel_stag = data.variables['uvel'][-1,0,:,:]
        self.vvel_stag = data.variables['uvel'][-1,0,:,:]

        try:
            self.wvel_ustag = data.variables['wvel_ho'][-1,0,:,:]
        except:
            self.wvel_ustag = data.variables['wvel'][-1,0,:,:]
        self.wvel_stag = (  self.wvel_ustag[1: , 1: ] + self.wvel_ustag[1: , :-1]
                          + self.wvel_ustag[:-1, :-1] + self.wvel_ustag[:-1, 1: ]) / 4.0

        self.uvel =  self.uvel_stag*math.cos(alpha) + self.wvel_stag*math.sin(alpha)
        self.vvel = -self.uvel_stag*math.sin(alpha) + self.wvel_stag*math.cos(alpha)

        self.x = (self.x0*math.cos(alpha)
                  + (self.usurf_stag[20,:]-7000.0)*math.sin(alpha)
                  )/1000.0 - 50.0
        self.y = self.y0/1000.0 - 50.0

def get_plot_data(test_file, bench_file, setup, config):
    test_plot_data = {}
    bench_plot_data = {}
    exp = config['name'].split('-')[-1]
    test_data = Dataset(test_file, 'r')
    bench_data = Dataset(bench_file, 'r')

    test = DataGrid(test_data)
    bench = DataGrid(bench_data)

    x_coord = setup['interp_points']
    y_coord = np.linspace(setup['y'][0], setup['y'][1], len(x_coord))

    test_plot_data['y_hat'] = y_coord
    test_plot_data['x_hat'] = x_coord
    bench_plot_data['y_hat'] = y_coord
    bench_plot_data['x_hat'] = x_coord

    if exp in ['a', 'c']:
        for var in config['interp_vars']:
            if var == 'usurf':
                # regular 2d linear interp. but faster.
                test2plot = interpolate.RectBivariateSpline(test.y_hat, test.x_hat,
                                                            test_data.variables[var][-1,:,:],
                                                            kx=1, ky=1, s=0)
                bench2plot = interpolate.RectBivariateSpline(bench.y_hat, bench.x_hat,
                                                             bench_data.variables[var][-1,:,:],
                                                             kx=1, ky=1, s=0)
            else:
                # regular 2d linear interp. but faster.
                test2plot = interpolate.RectBivariateSpline(test.y_hat, test.x_hat,
                                                            test_data.variables[var][-1,0,:,:],
                                                            kx=1, ky=1, s=0)
                bench2plot = interpolate.RectBivariateSpline(bench.y_hat, bench.x_hat,
                                                             bench_data.variables[var][-1,0,:,:],
                                                             kx=1, ky=1, s=0)

            test_plot_data[var] = test2plot(y_coord, x_coord, grid=False)
            bench_plot_data[var] = bench2plot(y_coord, x_coord, grid=False)

        test_plot_data['velnorm_extend'] = \
            np.linalg.norm(
                np.array([test_plot_data['uvel_extend'],
                          test_plot_data['vvel_extend'] ]),
                axis=0)
        bench_plot_data['velnorm_extend'] = \
            np.linalg.norm(
                np.array([bench_plot_data['uvel_extend'],
                          bench_plot_data['vvel_extend'] ]),
                axis=0)
    else:  # f
        alpha = math.radians(-3.0)

        test_rotated = RotatedGrid(alpha, test_data)
        bench_rotated = RotatedGrid(alpha, bench_data)

        for var in config['interp_vars']:
            # regular 2d linear interp. but faster.
            test2plot = interpolate.RectBivariateSpline(test_rotated.x, test_rotated.y,
                                                        getattr(test_rotated, var),
                                                        kx=1, ky=1, s=0)
            bench2plot = interpolate.RectBivariateSpline(bench_rotated.x, bench_rotated.y,
                                                         getattr(bench_rotated, var),
                                                         kx=1, ky=1, s=0)

            test_plot_data[var] = test2plot(y_coord, x_coord, grid=False)
            bench_plot_data[var] = bench2plot(y_coord, x_coord, grid=False)

        test_plot_data['velnorm'] = \
            np.linalg.norm(
                np.array([test_plot_data['uvel'],
                          test_plot_data['vvel'] ]),
                axis=0)
        bench_plot_data['velnorm'] = \
            np.linalg.norm(
                np.array([bench_plot_data['uvel'],
                          bench_plot_data['vvel'] ]),
                axis=0)

    return {'test': test_plot_data, 'bench': bench_plot_data}


"""
    Qunantum Transport simulator using QUEST.
    Copyright (C) 2014  K M Masum Habib <masum.habib@gmail.com>

    Last update: 05/12/2014
"""

import os
import pickle as pk
import numpy as np
import random as rn

from quest import setVerbosity, greet, vprint
from quest.vprint import nprint, dprint, eprint
from quest.linspace import linspace
from quest.atoms import AtomicStruct, SVec, LCoord
from quest.hamiltonian import TISurfKpParams4, TISurfKpParams, TI3DKpParams, GrapheneKpParams, GrapheneOneValleyKpParams, GrapheneTwoValleyKpParams, GrapheneTbParams, generateHamOvl
from quest.negf import CohRgfLoop
from quest.kpoints import KPoints
from quest.potential import LinearPot
from quest.utils import Timer, Workers, Quadrilateral, Point

class Transport(object):
    """
    Qunantum Transport simulator using QUEST.

    * Device geometry:
        ----------------------------------------
         ... |-1 | 0 | 1 |  ...  | N |N+1|N+2| ...
        ----------------------------------------
                   ^  <------^------>  ^
                 left      Device    right
               contact              contact

    Notes:
        1. Coherent RGF calculation for non-uniform devices requires at least
           5 blocks: 1 source block, 1 drain block and 3 device blocks as
           shown in the following fig:

           --------------------------
             | 0 | 1 | 2 | 3 | 4 |
           --------------------------
               ^  <--------->  ^
            source  Device   drain

           Block  0 and block 1 has to be identical.
           Simularly, block 3 and block 4 hast to be identical.

        2. Any uniform device needs atleast 3 blocks: two contact
           blocks and one device block.

    Attributes:

    """

    def __init__(self, workers = None):
        """
        Class constructor.
        param:
          workers: mpi.world - MPI world communicator.
        """

        self.verbosity = vprint.MSG_NORMAL # Verbosity level

        # MPI stuff
        if workers is None:
            self.workers = Workers()
        else:
            self.workers = Workers(workers)
        vprint.IAmMaster = self.workers.IAmMaster()

        # Timer
        self.clock = Timer()

        # output path settings
        self.OutPath    = "./out/"     # Output path
        self.OutFileName = "TR"        # Output file name prefix

        # Bias
        self.VDD        = np.zeros(1)  # Drain bias
        self.VGG        = np.zeros(1)  # Gate bias
        self.Vo         = 0.0          # Built in potential

        # Device geometry
        self.nb         = 11           # Length of the device+contacts
        self.nw         = 9            # Width of the device
        self.nh         = 1
        self.nbw        = []           # Width of all layers

        # for k-loop
        self.nk2         = 0            # number of k points along a2
        self.kp         = KPoints()    # K point generator

        # Hamiltonian type
        self.HAM_TI_SURF_KP  = 10       # TI surface k.p hamiltonian
        self.HAM_TI_SURF_KP4 = 11       # TI surface k.p hamiltonian with 4 spin basis set
        self.HAM_TI_3D_KP    = 15       # TI 3D k.p hamiltonian with 4 spin basis set
        self.HAM_GRAPHENE_KP = 20       # Graphene k.p Hamiltonian
        self.HAM_GRAPHENE_TWO_VALLEY_KP = 21       # Graphene k.p Hamiltonian
        self.HamType         = self.HAM_TI_SURF_KP

        # Device types
        self.COH_RGF_UNI    = 10       # Coherent uniform RGF type device
        self.COH_RGF_NON_UNI= 20       # Coherent Non-uniform RGF type device
        self.DevType        = self.COH_RGF_UNI

        # Potential profile type
        self.POT_LINEAR     = 10       # Linear potential model
        self.PotType        = self.POT_LINEAR

        # Calculations
        self.Calculations = {}
        self.Calculations["TE"] = 1

        #
        self.kT             = 0.0259        # Temperature in eV (300 K)
        self.ieta           = 1E-3j         # Contact imaginary potential
        self.mu             = 0.0           # Device Fermi level
        self.OrthoBasis     = True          # Orthogonal basis?
        self.Emin           =-1.0           # Minimum energy
        self.Emax           = 1.0           # Maximum energy
        self.dE             = 0.005         # Energy step
        self.AutoGenE       = False         # Generate grid automatically?

        # Debug stuffs
        self.DebugPotFile   = "dbg_pot.dat"
        self.DebugGeomFile  = "dbg_geom.svg"

        # Dry run
        self.DryRun         = False

        # Skip if resulting dat file already exists?
        self.SkipExistingSimulation = False

        # Print welcome message.
        nprint(greet())

    @property
    def verbosity(self):
        return self._verbosity
    @verbosity.setter
    def verbosity(self, value):
        self._verbosity = value
        # Set verbosity level of QUEST library.
        setVerbosity(self._verbosity)
        vprint.verbosity = self._verbosity

    def muS(self, VDD):
        """Calculates the fermi energy level of the source."""
        return self.mu - VDD*self.V.rVS

    def muD(self, VDD):
        """Calculates the drain energy level of the source."""
        return self.mu - VDD*self.V.rVD

    def VG(self, VGG, Vo, ig):
        """ Computes the gate voltage of gate #ig. """
        # Set gate voltages
        return VGG*self.V.rVG[ig] + Vo*self.V.rVo[ig]

    def addSource(self, sql, rVS = -0.5):
        """ Adds a source contact to the device for electrostatics calculation."""
        self.V.addSource(sql)
        self.V.rVS = rVS
        # just for pickling
        self.V.sql = sql

    def addDrain(self, dql, rVD = 0.5):
        """ Adds a drain contact to the device for electrostatics calculation."""
        self.V.addDrain(dql)
        self.V.rVD = rVD
        # just for pickling
        self.V.dql = dql

    def addGate(self, gql, rVo = 1, rVG = 1):
        """ Adds a gate to the device for electrostatics calculation."""
        self.V.addGate(gql)
        self.V.rVG.append(rVG)
        self.V.rVo.append(rVo)
         # just for pickling
        self.V.gql.append(gql)

    def addLinearRegion(self, lql):
        """ Adds a linear region to the device for electrostatics calculation."""
        self.V.addLinearRegion(lql)
        # just for pickling
        self.V.lql.append(lql)

    def createAtomicGeom(self):
        """ Creates atomistic geometry. """

        nprint("\n Creating atomistic geometry ...")

        if not hasattr(self, "hp"): # if hp does not exist, create it.
            if self.HamType == self.HAM_TI_SURF_KP:
                self.hp = TISurfKpParams()
            elif self.HamType == self.HAM_TI_SURF_KP4:
                self.hp = TISurfKpParams4()
            elif self.HamType == self.HAM_TI_3D_KP:
                self.hp = TI3DKpParams()
            elif self.HamType == self.HAM_GRAPHENE_KP:
                self.hp = GrapheneOneKpParams()
            elif self.HamType == self.HAM_GRAPHENE_TWO_VALLEY_KP:
                self.hp = GrapheneTwoKpParams()
            else:
                raise RuntimeError(" Unsupported Hamiltonian type. ")

        dev = AtomicStruct(self.hp.ptable)
        dev.genSimpleCubicStruct(self.hp.ptable[0], self.hp.a, self.nb-2, self.nw, self.nh)

        if not hasattr(self, "nlc"):
            self.nlc = self.nw

        if not hasattr(self, "nrc"):
            self.nrc = self.nw

        cont_left = AtomicStruct(self.hp.ptable)
        cont_left.genSimpleCubicStruct(self.hp.ptable[0], self.hp.a, 1, self.nlc, self.nh)
        xs = dev.xmin
        cont_left += np.array([xs, 0, 0])
        cont_left += LCoord(-1, 0, 0)

        cont_right = AtomicStruct(self.hp.ptable)
        cont_right.genSimpleCubicStruct(self.hp.ptable[0], self.hp.a, 1, self.nrc, self.nh)
        xs = dev.xmax
        cont_right += np.array([xs, 0, 0])
        cont_right += LCoord(1, 0, 0)

        self.geom = cont_left + dev + cont_right
        self.lyr_0 = cont_left
        self.lyr_nbm1 = cont_right
        cont_left += LCoord(-1, 0, 0)
        cont_right += LCoord(1, 0, 0)
        self.lyr_0m1 = cont_left
        self.lyr_nb = cont_right

        self.nbw = [self.nw]*self.nb
        self.nbw[0] = self.nlc
        self.nbw[self.nb-1] = self.nrc
        self.updateBoundingBox()

        nprint(" done.")

    def createRoughEdges(self, sigma):
        """
        Creates rough edges. Works only for sorted lattice points.
        For rectangular lattice
        """

        nprint("\n Creating rough edges ...")

        if (self.HamType == self.HAM_TI_SURF_KP
                    or self.HamType == self.HAM_TI_SURF_KP4
                    or self.HamType == self.HAM_GRAPHENE_KP
                    or self.HamType == self.HAM_GRAPHENE_TWO_VALLEY_KP):
            nw = []
            geom = AtomicStruct()
            beg = 0
            # loop through the layers and
            # remove some atoms randomly from the edges.
            for ib in range(0, self.nb):
                end = beg + self.nw - 1
                lyr = self.geom.span(beg, end)
                beg = end + 1

                if ( ib > 1 and ib < self.nb - 2):
                    # Remove or add some atoms from the bottom edge.
                    nr = int(rn.gauss(0, sigma))
                    while abs(nr) > 3*sigma:
                        nr = int(rn.gauss(0, sigma))
                    if nr <= 0:
                        lyr = lyr.span(abs(nr), lyr.NumOfAtoms - 1)
                    else:
                        xtra = lyr.span(0, nr-1)
                        lv = SVec(0, - self.hp.ay*nr, 0)
                        xtra = xtra + lv
                        lyr = xtra + lyr

                    # Remove or add some atoms from the top edge
                    nr = int(rn.gauss(0, sigma))
                    while abs(nr) > 3*sigma:
                        nr = int(rn.gauss(0, sigma))
                    if nr <= 0:
                        lyr = lyr.span(0, lyr.NumOfAtoms - 1 - abs(nr))
                    else:
                        xtra = lyr.span(lyr.NumOfAtoms - nr, lyr.NumOfAtoms - 1)
                        lv = SVec(0, self.hp.ay*nr, 0)
                        xtra = xtra + lv
                        lyr = lyr + xtra

                # save this layer
                geom = geom + lyr
                nw.append(lyr.NumOfAtoms)

            # save geometry
            self.geom = geom
            self.nbw = nw
            self.DevType = self.COH_RGF_NON_UNI

            # update bounding box
            self.updateBoundingBox()

        else:
            raise RuntimeError(" Unsupported Hamiltonian type. ")

        nprint(" done.")

    def generateHamiltonian(self):
        """ Generates hamiltonian and overlap matrices. """

        nprint("\n Generating Hamiltonian matrix ...")

        # For uniform RGF blocks
        if (self.DevType == self.COH_RGF_UNI):
            # no k-loop, real space hamiltonian only
            if self.kp.N == 0:

                lyr0 = self.geom.span(0, self.nw*self.nh-1)               # extract block # 0
                lyr1 = self.geom.span(self.nw*self.nh, 2*self.nw*self.nh-1)    # extract block # 1

                self.H0, self.S0 = generateHamOvl(self.hp, lyr0, lyr0)

                self.Hl, S = generateHamOvl(self.hp, lyr1, lyr0)

#                np.set_printoptions(linewidth=200)
#                print "\nS0\n"
#                print self.S0
#                print "\nH0\n"
#                print self.H0
#                print "\nHl\n"
#                print self.Hl
#                self.geom.exportGjf('dbg_geom.gjf')
#                lyr0.exportGjf('lyr0.gjf')
#                lyr1.exportGjf('lyr1.gjf')
#                lyr01 = lyr0+lyr1;
#                lyr01.exportGjf('lyr01.gjf')

            # nearest neighbors in transverse direction for k-loop
            else:
                self.H0 = []
                self.S0 = []
                self.Hl = []
                self.pv = []
                self.pvl = []

                lv = self.geom.LatticeVector
                self.pv.append(lv*LCoord(0,0,0))
                self.pvl.append(lv*LCoord(0,0,0))
                lyr0 = self.geom.span(0, self.nw*self.nh-1)               # extract block # 0
                lyr1 = self.geom.span(self.nw*self.nh, 2*self.nw*self.nh-1)    # extract block # 1

                H, S = generateHamOvl(self.hp, lyr0, lyr0)
                self.H0.append(H); self.S0.append(S)
                H, S = generateHamOvl(self.hp, lyr1, lyr0)
                self.Hl.append(H)

                self.pv.append(lv*LCoord(0,1,0))
                lyr0top = lyr0 + self.pv[1]                # top neighbor of layer 0
                H, S = generateHamOvl(self.hp, lyr0, lyr0top)
                self.H0.append(H)

                self.pv.append(lv*LCoord(0, -1, 0))
                lyr0bot = lyr0 + self.pv[2]                # bottom neighbor of layer 0
                H, S = generateHamOvl(self.hp, lyr0, lyr0bot)
                self.H0.append(H)

                self.pvl.append(lv*LCoord(-1,1,0))
                H, S = generateHamOvl(self.hp, lyr1, lyr0top)
                self.Hl.append(H)

                self.pvl.append(lv*LCoord(-1,-1,0))
                H, S = generateHamOvl(self.hp, lyr1, lyr0bot)
                self.Hl.append(H)
        # Non-uniform RGF blocks
        elif (self.DevType == self.COH_RGF_NON_UNI):
            self.H0 = []
            self.S0 = []
            self.Hl = []
            beg = 0
            for ib in range(0, self.nb):                    # setup the block hamiltonian
                end = beg + self.nbw[ib] - 1

                # generate H_i,i and S_i,i
                lyri = self.geom.span(beg, end)            # extract block # i
                H0,S0 = generateHamOvl(self.hp, lyri, lyri)
                self.H0.append(H0)
                self.S0.append(S0)

                # generate H_i,i-1
                if ib > 0:
                    Hl,Sl = generateHamOvl(self.hp, lyri, lyrim1)
                    self.Hl.append(Hl)

                lyrim1 = lyri
                beg = end + 1
            # Coupling matrix between two blocks of left contact
            Hl,Sl = generateHamOvl(self.hp, self.lyr_0, self.lyr_0m1)
            self.Hl.insert(0, Hl)
            # Coupling matrix between two blocks of right contact
            Hl,Sl = generateHamOvl(self.hp, self.lyr_nb, self.lyr_nbm1)
            self.Hl.append(Hl)

        nprint(" done.")

    def setupPotential(self):
        """ Sets up the potential profile """

        nprint("\n Setting up potential ...")
        # Linear
        if (self.PotType == self.POT_LINEAR):
            self.V = LinearPot(self.geom)
        nprint(" done.")

    def check(self):
        ret = True
        if (self.DevType == self.COH_RGF_UNI):
            ret = ret and (self.nb >= 3)
        if (self.DevType == self.COH_RGF_NON_UNI):
            ret = ret and (self.nb >= 5)

        return ret

    def runBiasStep(self, VGG, Vo, VDD):
        """Runs the sumulation."""

        # Set drain and Fermi levels
        self.rgf.mu(self.muD(VDD), self.muS(VDD))

        nprint("\n Bias loop:")
        nprint("\n  VGG = " + str(VGG) + ", Vo = " + str(Vo)
                    + ", VDD = " + str(VDD) + ".")

        fileName = self.OutFileName + "_VGG{0:2.3f}_Vo{1:2.3f}_VDD{2:2.3f}".format(VGG, Vo, VDD)

        # skip calculation if result file exists.
        if self.SkipExistingSimulation == True:
            if os.path.isfile(self.OutPath + fileName + ".dat"):
                nprint("\n  Result exists, skipping.")
                return

        # Set gate voltages
        for ig in range(self.V.NG):
            VG = self.VG(VGG, Vo, ig)
            # bias voltage for gate # ig
            self.V.VG(ig, VG)
            # For linear potential profile
            if (self.PotType == self.POT_LINEAR):
                # set potential of source
                if (ig == 0):
                    self.V.VS(VG)
                # set potential of drain
                if (ig == self.V.NG - 1):
                    self.V.VD(VG)

        # Potential of linear region
        if (self.PotType == self.POT_LINEAR):
            for il in range(self.V.NLR):
                VG1 = self.VG(VGG, Vo, il)
                VG2 = self.VG(VGG, Vo, il+1)
                self.V.VLR(il, VG1, VG2)

        # Print some useful information
        nprint("\n")
        nprint(self.dynamicnstr())

        # Compute the electrostatic potential
        self.V.compute()
        eprint(self.V)

        # Save the calculated potential.
        if (self.verbosity >= vprint.MSG_DEBUG):
            if (self.workers.IAmMaster()):
                # Create directory if not exist.
                if not os.path.exists(self.OutPath):
                    os.makedirs(self.OutPath)
                self.V.exportPotential(self.OutPath + self.DebugPotFile)

        # Export potential to NEGF
        beg = 0
        self.Vo = []
        for ib in range(self.nb):                  # setup the block hamiltonian
            end = beg + self.nbw[ib]*self.nh - 1
            self.Vo.append(self.V.toOrbPot(beg, end))
            self.rgf.V(self.Vo[ib], ib)
            beg = end + 1

        # Create energy grid
        if (self.AutoGenE):
            Emin = self.muD(VDD) - 10*self.kT
            Emax = self.muS(VDD) + 10*self.kT
        else:
            Emin = self.Emin
            Emax = self.Emax

        EE = linspace(Emin, Emax, self.dE)

        nprint("\n Energy grid:\n  Min: " + str(Emin) + ", max: "
                    + str(Emax) + ", interval " + str(self.dE)
                    + ".")
        nprint("\n Total " + str(len(EE)) + " energy point(s) "
                    + "running on " + str(self.workers.N()) + " CPU(s): "
                    + str(int(round(len(EE)/self.workers.N()))) + " pts/CPU ... \n")

        # Set energy
        self.rgf.E(EE)

        # Run the simulation
        if (self.DryRun == False):
            self.rgf.run()
        nprint("\n  done.")

        nprint("\n Saving results to disk ...")
        # save results
        if (self.workers.IAmMaster()):
            # Create directory if not exist.
            if not os.path.exists(self.OutPath):
                os.makedirs(self.OutPath)
            # save results to file.
            if (self.DryRun == False):
                fo = open(self.OutPath + fileName + ".dat", "wt")
                fo.close()
                self.rgf.save(self.OutPath + fileName + ".dat")

        nprint(" done.\n")
        nprint(" ------------------------------------------------------------------")

    def run(self):
        """Runs the sumulation."""

        self.clock.tic()

        nprint("\n\n Starting simulation ...")

        # Print simulation info
        nprint(self.staticnstr())
        eprint(self.debugstr())

        if (self.check() == False):
            raise error(" ERROR: Check failed !!!")

        # Debug print
        dprint("\n" + str(self.VDD))
        dprint("\n" + str(self.VGG))

        # save simulation parameters
        if (self.workers.IAmMaster()):
            # Create directory if not exist.
            if not os.path.exists(self.OutPath):
                os.makedirs(self.OutPath)
            # save simulation parameters to a pickle file
            fp = open(self.OutPath + self.OutFileName + ".pkl", 'wb')
            pk.dump(self, fp)
            fp.close()

        # Save electrostatic geometry
        if (self.verbosity >= vprint.MSG_DEBUG):
            if (self.workers.IAmMaster()):
                # Create directory if not exist.
                if not os.path.exists(self.OutPath):
                    os.makedirs(self.OutPath)
                self.V.exportSvg(self.OutPath + self.DebugGeomFile)

        # Configure the RGF solver
        if self.kp.N == 0: # without k-loop
            self.rgf = CohRgfLoop(self.workers, self.nb, self.kT, self.ieta,
                    self.OrthoBasis)
        else:# with k-loop
            self.rgf = CohRgfLoop(self.workers, self.nb, self.kT, self.ieta,
                    self.OrthoBasis, 2)

        # Setup H and S
        if (self.DevType == self.COH_RGF_UNI):        # for uniform RGF blocks

            for ib in range(0, self.nb+1):            # setup the block hamiltonian
                if self.kp.N == 0: # no k-loop
                    if (ib != self.nb):
                        self.rgf.H0(self.H0, ib)          # H0: 0 to N+1
                        self.rgf.S0(self.S0, ib)          # S0: 0 to N+1
                    self.rgf.Hl(self.Hl, ib)              # Hl: 0 to N+2

                else: # we have k-loop, add transverse neighbors
                    self.rgf.k(self.kp.kp)                      # set k-points
                    if (ib != self.nb):
                        self.rgf.H0(self.H0[0], ib, 0)          # H0_i,i: 0 to N+1
                        self.rgf.H0(self.H0[1], ib, 1)          # H0_i,i+1: 0 to N+1
                        self.rgf.H0(self.H0[2], ib, 2)          # H0_i,i-1: 0 to N+1
                        self.rgf.S0(self.S0[0], ib, 0)          # S0_i,i: 0 to N+1
                        self.rgf.pv0(self.pv[0], ib, 0)
                        self.rgf.pv0(self.pv[1], ib, 1)
                        self.rgf.pv0(self.pv[2], ib, 2)
                    self.rgf.Hl(self.Hl[0], ib, 0)              # Hl_i,i: 0 to N+2
                    self.rgf.Hl(self.Hl[1], ib, 1)              # Hl_i,i: 0 to N+2
                    self.rgf.Hl(self.Hl[2], ib, 2)              # Hl_i,i: 0 to N+2
                    self.rgf.pvl(self.pvl[0], ib, 0)
                    self.rgf.pvl(self.pvl[1], ib, 1)
                    self.rgf.pvl(self.pvl[2], ib, 2)

        elif (self.DevType == self.COH_RGF_NON_UNI):  # Non-uniform RGF blocks
            for ib in range(0, self.nb):                 # setup the block hamiltonian
                self.rgf.H0(self.H0[ib], ib)             # H0: 0 to N+1=nb-1
                self.rgf.S0(self.S0[ib], ib)              # S0: just the identity matrix stored in S0(0)
                if ib > 0:
                    self.rgf.Hl(self.Hl[ib], ib)     # Hl: 1 to N+1=nb-1
            self.rgf.Hl(self.Hl[0], 0)               # Hl(0) = H_0,-1
            self.rgf.Hl(self.Hl[ib+1], ib+1)           # Set H_N+2,N+1 = H_N+1,N
        # Clean up unused memory; we alredy have copies of these variables
        # in self.rgf.
        self.H0 = None
        self.Hl = None
        self.S0 = None
        self.Sl = None

        # Enable calculations
        for type, value in self.Calculations.items():
            # transmission
            if (type == "TE"):
                self.rgf.enableTE(value)
            # current
            if (type == "I"):
                if (isinstance( value, int)):
                    self.rgf.enableI(value, 0, 1)
                else:
                    for I in self.Calculations["I"]:
                        if ("Block" in I and I["Block"] == "All"):
                            for ib in range(0, self.nb-1):
                                self.rgf.enableI(I["N"], ib, ib+1)
                        else:
                            self.rgf.enableI(I["N"], I["From"], I["To"])
            if (type == "DOS"):
                self.rgf.enableDOS(value)
            if (type == "n"):
                if (isinstance( value, int)):
                    self.rgf.enablen(value)
                else:
                    for n in self.Calculations["n"]:
                        if (n["Block"] == "All"):
                            for ib in range(1, self.nb-1):
                                self.rgf.enablen(n["N"], ib)
                        else:
                            self.rgf.enablen(n["N"], n["Block"])
        if hasattr(self, "atomsTracedOver"):
            self.rgf.atomsTracedOver(self.atomsTracedOver);

        # Loop over drain and gate bias
        for VDD in self.VDD:
            for VGG in self.VGG:
                self.runBiasStep(VGG, self.Vo, VDD)
                pass
        self.clock.toc()
        nprint("\n" + str(self.clock) + "\n")

    def updateBoundingBox(self):
        """Updates the bounding box of atomistic geometry."""
        # Just to make sure that no point of gate regions is
        # at the border
        a = self.hp.a
        delta = a*7.0/220.0
        self.xmn = self.geom.xmin - delta
        self.xmx = self.geom.xmax + delta
        self.ymn = self.geom.ymin - delta
        self.ymx = self.geom.ymax + delta

    def __str__(self):
        msg = self.staticnstr()
        msg += self.dynamicnstr()
        return msg

    def staticnstr(self):
        msg = "\n Transport simulation parameters:"

        msg += "\n  Device geometry:"
        msg += "\n  ---------------------------------------------"
        msg += "\n  ... |-1 | 0 | 1 |   ...  | " + str(self.nb-2) + " | " + str(self.nb-1) + " | " + str(self.nb) + " | ..."
        msg += "\n  ---------------------------------------------"
        msg += "\n            ^  <-------^------->  ^"
        msg += "\n          left       Device     right"
        msg += "\n        contact                contact"

        # Bias information
        msg += "\n Bias: "
        msg += "\n  VDD: min " + str(min(self.VDD)) + ", max " + str(max(self.VDD))
        msg += ", number " + str(len(self.VDD))
        msg += "\n  VGG: min " + str(min(self.VGG)) + ", max " + str(max(self.VGG))
        msg += ", number " + str(len(self.VGG))
        msg += "\n  Vo: " + str(self.Vo)

        # Device information
        msg += "\n Device:"
        msg += "\n  Lengh: " + str(self.geom.xl) + "(" + str(self.geom.xl/10.0) + " nm)"
        msg += ", Width: " +  str(self.geom.yl) + "(" + str(self.geom.yl/10.0) + " nm)"
        msg += ", Height: " + str(self.geom.zl)+ "(" + str(self.geom.zl/10.0) + " nm)"
        msg += "\n  Num of atoms: " + str(self.geom.NumOfAtoms)
        msg += ", Num of orbitals: " +  str(self.geom.NumOfOrbitals)

        # Print Hamiltonian parameters.
        msg += "\n"
        msg += str(self.hp)

        # Calculations to perform
        msg += " Calculations:\n"
        for type, value in self.Calculations.items():
            if (type == "TE"):
                msg += "  Transmission.\n"
            if (type == "I"):
                if (isinstance( value, int)):
                    msg += "  Current at left contact"
                    msg += " (" + str(I["N"]) + ").\n"
                else:
                    for I in self.Calculations["I"]:
                        if ("Block" in I and I["Block"] == "All"):
                            msg += "  Current of all blocks"
                            msg += " (" + str(I["N"]) + ").\n"
                        else:
                            msg += "  Current from block # " + str(I["From"])
                            msg += " to block # " + str(I["To"])
                            msg += " (" + str(I["N"]) + ").\n"
            if (type == "n"):
                if (isinstance( value, int)):
                    msg += "  Electron density of the device"
                    msg += " (" + str(n["N"]) + ").\n"
                else:
                    for n in self.Calculations["n"]:
                        if (n["Block"] == "All"):
                            msg += "  Electron density of all blocks"
                            msg += " (" + str(n["N"]) + ").\n"
                        else:
                            msg += "  Electron density of block # " + str(n["Block"])
                            msg += " (" + str(n["N"]) + ").\n"

        msg += "  Save output at: " + self.OutPath + self.OutFileName + "*\n"

        # End of info section
        msg += " ------------------------------------------------------------------"

        return msg

    def dynamicnstr(self):
        # NEGF parameters
        msg = str(self.rgf)
        # Electrostatics
        msg += "\n"
        msg += str(self.V)

        return msg

    def debugstr(self):
        msg = " Debugging information: \n"
        msg = "  Atomic structure: \n"
        msg += str(self.geom)
        msg += " ------------------------------------------------------------------"
        return msg

    def __getstate__(self):
        dct = dict(self.__dict__)
        del dct['workers']
        del dct['clock']
        del dct['kp']
        #del dct['H0']
        #del dct['Hl']
        #del dct['S0']
        #del dct['Sl']
        #del dct['geom']
        #del dct['rgf']
        return dct

    def __setstate__(self, dct):
        self.__dict__.update(dct)

"""
 Loads and returns a Transport object from pickle file
"""
def loadTransport(fileName):
    pf = open(fileName)
    tr = pk.load(pf)
    pf.close()
    return tr


#!/usr/bin/python3

"""Read .out files from ORCA IRC calculations and create graphs."""

import argparse

import matplotlib.pyplot as plt
import numpy as np
from scipy.constants import calorie
from scipy.constants import kilo
from scipy.constants import N_A
from scipy.constants import physical_constants
from scipy import interpolate

from rmsd import calc_rmsd
from rmsd import read_xyz

hartree, _, _ = physical_constants["Hartree energy"]

def main():
    """Run main procedure."""
    parser = argparse.ArgumentParser()
    parser.add_argument("out_file", type=argparse.FileType("r"), default="-")
    parser.add_argument("--classic", action="store_true")
    args = parser.parse_args()

    irc = False
    data = []
    for line in args.out_file:
        if "Step        E(Eh)      dE(kcal/mol)  max(|G|)   RMS(G)" in line:
            irc = True
        elif irc:
            line = line.strip()
            if line:
                fields = line.split()
                data.append([float(x) for x in fields[:5]])
                if len(fields) > 5 and fields[5] == "<=":
                    ts_step = int(data[-1][0]) - 1  # first step is one
            else:
                break

    try:
        coords = read_xyz(args.out_file.name.replace(".out", "_TSOpt_IRC_Full_trj.xyz"))[2]
    except ValueError:
        coords = read_xyz(args.out_file.name.replace(".out", "_IRC_Full_trj.xyz"))[2]

    rmsd = [0.0]
    for i in range(1, len(coords)):
        rmsd.append(rmsd[-1] + calc_rmsd(coords[i - 1], coords[i]))
    rmsd = np.array(rmsd)

    data = np.array(data)
    # xi = (data[:, 0] - data[:, 0].min()) / data[:, 0].max()
    xi = rmsd
    y = data[:, 1] - data[:, 1].min()

    forward_barrier = y.max() - y[0]
    backward_barrier = y.max() - y[-1]
    print(
        f"forward barrier  = {forward_barrier:6.4f} Eh = {forward_barrier * hartree * N_A / kilo:5.1f} kJ/mol = {forward_barrier * hartree * N_A / (kilo * calorie):5.1f} kcal/mol"
    )
    print(
        f"backward barrier = {backward_barrier:6.4f} Eh = {backward_barrier * hartree * N_A / kilo:5.1f} kJ/mol = {backward_barrier * hartree * N_A / (kilo * calorie):5.1f} kcal/mol"
    )

    if not args.classic:
        xi_new = np.linspace(xi.min(), xi.max(), 10000)

        # points = ~np.isclose(xi, xi[ts_step])
        # f = interpolate.InterpolatedUnivariateSpline(xi[points], y[points])
        f = interpolate.InterpolatedUnivariateSpline(xi, y, k=4)
        fp = f.derivative()
        fpp = f.derivative(n=2)

        plt.subplot(311)
    else:
        plt.subplot(211)

    plt.plot(xi, y, "o", label="calculated")
    plt.vlines(xi[ts_step], y.min(), y.max())
    plt.xlabel(r"IRC ($\xi$)")
    plt.ylabel(r"Potential Energy, V($\xi$) [Eh]")

    if not args.classic:
        # add interpolation
        pass
        plt.plot(xi_new, f(xi_new), "--", label="interpolated")
        plt.legend()

        plt.subplot(312)
        rf = -fp(xi_new)
        plt.plot(xi_new, rf, "--")
        plt.vlines(xi[ts_step], rf.min(), rf.max())
        plt.xlabel(r"IRC ($\xi$)")
        plt.ylabel(r"Reaction Force, F($\xi$) [Eh/$\Delta\xi$]")

        # TODO(schneiderfelipe): get max and min force along the coordinate
        # and do the usual analysis
        plt.subplot(313)
        rfc = fpp(xi_new)
        plt.plot(xi_new, rfc, "--")
        plt.vlines(xi[ts_step], rfc.min(), rfc.max())
        plt.xlabel(r"IRC ($\xi$)")
        plt.ylabel(r"Reaction Force Constant, $\kappa$($\xi$) [Eh/$\Delta\xi^2$]")
    else:
        plt.subplot(212)
        plt.plot(xi, data[:, 3], "o", label="max(|G|)")
        plt.vlines(xi[ts_step], 0, data[:, 3].max())
        plt.plot(xi, data[:, 4], "o", label="RMS(G)")
        plt.xlabel(r"IRC ($\xi$)")
        plt.ylabel(r"Gradient, G($\xi$) [Eh/Bohr]")
        plt.legend()

    # plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()