import os
import shutil

import numpy as np
import tensorflow as tf
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.layers import concatenate

from seisnn.core import Instance
from seisnn.model.generator import nest_net
from seisnn.model.attention import transformer
from seisnn.model.GAN_model import build_discriminator, build_cgan, build_patch_discriminator
import seisnn.example_proto
import seisnn.io
import seisnn.sql
import seisnn.utils

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

class BaseTrainer:
    @staticmethod
    def get_dataset_length(database=None, tfr_list=None):
        count = None
        try:
            db = seisnn.sql.Client(database)
            tfr_list = seisnn.utils.flatten_list(tfr_list)
            counts = db.get_tfrecord(path=tfr_list, column='count')
            count = sum(seisnn.utils.flatten_list(counts))
        except Exception as error:
            print(f'{type(error).__name__}: {error}')

        return count

    @staticmethod
    def get_model_dir(model_instance, remove=False):
        config = seisnn.utils.Config()
        save_model_path = os.path.join(config.models, model_instance)

        if remove:
            shutil.rmtree(save_model_path, ignore_errors=True)
        seisnn.utils.make_dirs(save_model_path)

        save_history_path = os.path.join(save_model_path, "history")
        seisnn.utils.make_dirs(save_history_path)

        return save_model_path, save_history_path

class GeneratorTrainer(BaseTrainer):
    """
    Trainer class.
    """

    def __init__(self, database=None, model=nest_net(),
                 optimizer=tf.keras.optimizers.Adam(1e-4),
                 loss=tf.keras.losses.BinaryCrossentropy()):
        """
        Initialize the trainer.

        :param database: SQL database.
        :param model: keras model.
        :param optimizer: keras optimizer.
        :param loss: keras loss.
        """
        self.generator_model = transformer(img_rows=1, img_cols=3008,
                                           color_type=3,
                                           num_class=3)
        self.discriminator_model = build_discriminator(img_rows=1,
                                                       img_cols=3008,
                                                       color_type=3,
                                                       num_class=3)
        self.generator_optimizer = Adam(learning_rate=1e-3)
        self.generator_model.compile(loss='binary_crossentropy',
                                     optimizer=self.generator_optimizer)

        self.discriminator_model.trainable = False

        self.cgan_model = build_cgan(self.generator_model,
                                     self.discriminator_model)

        loss = [tf.keras.losses.BinaryCrossentropy(),
                tf.keras.losses.BinaryCrossentropy()]
        loss_weights = [100000, 1]
        self.cgan_optimizer = Adam(learning_rate=1e-3)
        self.cgan_model.compile(loss=loss, loss_weights=loss_weights,
                                optimizer=self.cgan_optimizer)
        self.discriminator_optimizer = Adam(learning_rate=1e-3)
        self.discriminator_model.trainable = True
        self.discriminator_model.compile(loss='binary_crossentropy',
                                         optimizer=self.discriminator_optimizer)

        self.database = database
        self.model = self.cgan_model
        self.optimizer = optimizer

    def train_loop(self,
                   tfr_list, model_name,
                   epochs=1, batch_size=1,
                   log_step=100, plot=False,
                   remove=False):
        """
        Main training loop.

        :param tfr_list: List of TFRecord path.
        :param str model_name: Model directory name.
        :param int epochs: Epoch number.
        :param int batch_size: Batch size.
        :param int log_step: Logging step interval.
        :param bool plot: Plot training validation,
            False save fig, True show fig.
        :param bool remove: If True, remove model folder before training.
        :return:
        """
        model_path, history_path = self.get_model_dir(model_name,
                                                      remove=remove)

        ckpt = tf.train.Checkpoint(
            generator_model=self.generator_model,
            discriminator_model=self.discriminator_model,
            cgan_model=self.cgan_model,
            generator_optimizer=self.generator_optimizer,
            discriminator_optimizer=self.discriminator_optimizer,
            cgan_optimizer=self.cgan_optimizer,
        )
        ckpt_manager = tf.train.CheckpointManager(ckpt, model_path,
                                                  max_to_keep=100)

        if ckpt_manager.latest_checkpoint:
            ckpt.restore(ckpt_manager.latest_checkpoint)
            last_epoch = len(ckpt_manager.checkpoints)
            print(f'Latest checkpoint epoch {last_epoch} restored!!')
        dataset = seisnn.io.read_dataset(tfr_list)
        dataset = dataset.shuffle(5000)
        val = next(iter(dataset.batch(1)))
        metrics_names = ['loss', 'val']

        # data_len = self.get_dataset_length(self.database,tfr_list)

        for epoch in range(epochs):
            print(f'epoch {epoch + 1} / {epochs}')

            n = 0
            progbar = tf.keras.utils.Progbar(
                190000, stateful_metrics=metrics_names)

            for train in dataset.prefetch(100).batch(batch_size):
                d_loss, g_loss = self.train_step(train, val)
                values = [('d_loss', d_loss),
                          ('g_loss', g_loss)]
                progbar.add(len(train['id']), values=values)

                n += 1
                if n % log_step == 0:
                    val['predict'] = self.generator_model(val['trace'])
                    concate = concatenate([val['trace'], val['predict']],
                                          axis=3)
                    score = self.discriminator_model(concate)
                    print(score)
                    title = f'epoch{epoch + 1:0>2}_step{n:0>5}___'

                    val['id'] = tf.convert_to_tensor(
                        title.encode('utf-8'), dtype=tf.string)[tf.newaxis]

                    example = next(seisnn.example_proto.batch_iterator(val))
                    instance = Instance(example)

                    if plot:
                        instance.plot()
                    else:
                        instance.plot(save_dir=history_path)

            ckpt_save_path = ckpt_manager.save()
            print(f'Saving checkpoint to {ckpt_save_path}')
            self.generator_model.save(f'/home/andy/Models/{model_name}.h5')

    def train_step(self, train, val):
        """
        Training step.

        :param train: Training data.
        :param val: Validation data.
        :rtype: float
        :return: predict loss, validation loss
        """

        real = np.ones((train['trace'].shape[0], 1))
        fake = np.zeros((train['trace'].shape[0], 1))
        g_pred = self.generator_model(train['trace'], training=False)
        concat = concatenate((train['trace'], g_pred), axis=3)
        f_disc_loss = self.discriminator_model.train_on_batch(concat, fake)
        concat = concatenate((train['trace'], train['label']), axis=3)
        r_disc_loss = self.discriminator_model.train_on_batch(concat, real)
        disc_loss = 0.5 * (f_disc_loss + r_disc_loss)
        self.discriminator_model.trainable = False
        gen_loss = self.cgan_model.train_on_batch(train['trace'],
                                                  [train['label'], real])
        self.discriminator_model.trainable = True

        return np.array(disc_loss), np.array(gen_loss)

if __name__ == "__main__":
    database = 'CWB.db'
    db = seisnn.sql.Client(database)
    tfr_list = db.get_tfrecord(from_date='2000-01-01', to_date='2018-12-31', column='path')

    model_instance = 'qqq'
    trainer = GeneratorTrainer(database)
    trainer.train_loop(tfr_list, model_instance,
                       batch_size=250, epochs=100,
                       plot=True)


from app import db
from carculator import *
import json
import itertools
import csv
from app.models import Task
import numpy as np

class Calculation:
    def __init__(self):

        bs = BackgroundSystemModel()
        self.electricity_mix = bs.electricity_mix
        self.biogasoline = bs.biogasoline
        self.biodiesel = bs.biodiesel
        self.biomethane = bs.biomethane
        self.region_map = bs.region_map
        self.cip = CarInputParameters()
        self.cip.static()
        self.d_categories = {
            self.cip.metadata[a]["name"]: self.cip.metadata[a]["category"]
            for a in self.cip.metadata
        }
        self.dcts, self.arr = fill_xarray_from_input_parameters(self.cip)
        self.d_pt_en = {
            "Petrol": "ICEV-p",
            "Diesel": "ICEV-d",
            "CNG": "ICEV-g",
            "Electric": "BEV",
            "Fuel cell": "FCEV",
            "Hybrid-petrol": "HEV-p",
            "Hybrid-diesel": "HEV-d",
            "(Plugin) Hybrid-petrol": "PHEV-p",
            "(Plugin) Hybrid-diesel": "PHEV-d",
        }
        self.d_pt_it = {
            "Benzina": "ICEV-p",
            "Diesel": "ICEV-d",
            "Gas compresso": "ICEV-g",
            "Elettrica": "BEV",
            "Cella a combustibile": "FCEV",
            "Ibrido benzina": "HEV-p",
            "Ibrido diesel": "HEV-d",
            "Ibrido-benzina (Plugin)": "PHEV-p",
            "Ibrido-diesel (Plugin)": "PHEV-d",
        }
        self.d_pt_de = {
            "Benzin": "ICEV-p",
            "Diesel": "ICEV-d",
            "Komprimiertes Gas": "ICEV-g",
            "Elektrisch": "BEV",
            "Brennstoffzelle": "FCEV",
            "Hybrid-Benzin": "HEV-p",
            "Hybrid-Diesel": "HEV-d",
            "(Plugin) Hybrid-Benzin": "PHEV-p",
            "(Plugin) Hybrid-Diesel": "PHEV-d",
        }
        self.d_pt_fr = {
            "Essence": "ICEV-p",
            "Diesel": "ICEV-d",
            "Gaz comprimé": "ICEV-g",
            "Electrique": "BEV",
            "Pile à combustible": "FCEV",
            "Hybride-essence": "HEV-p",
            "Hybride-diesel": "HEV-d",
            "Hybride-essence rechargeable": "PHEV-p",
            "Hybride-diesel rechargeable": "PHEV-d",
        }
        self.d_pt_all = {
            "Petrol": "ICEV-p",
            "Diesel": "ICEV-d",
            "CNG": "ICEV-g",
            "Electric": "BEV",
            "Fuel cell": "FCEV",
            "Hybrid-petrol": "HEV-p",
            "Hybrid-diesel": "HEV-d",
            "(Plugin) Hybrid-petrol": "PHEV-p",
            "(Plugin) Hybrid-diesel": "PHEV-d",
            "Benzina": "ICEV-p",
            "Gas compresso": "ICEV-g",
            "Elettrica": "BEV",
            "Cella a combustibile": "FCEV",
            "Ibrido benzina": "HEV-p",
            "Ibrido diesel": "HEV-d",
            "Ibrido-benzina (Plugin)": "PHEV-p",
            "Ibrido-diesel (Plugin)": "PHEV-d",
            "Benzin": "ICEV-p",
            "Komprimiertes Gas": "ICEV-g",
            "Elektrisch": "BEV",
            "Brennstoffzelle": "FCEV",
            "Hybrid-Benzin": "HEV-p",
            "Hybrid-Diesel": "HEV-d",
            "(Plugin) Hybrid-Benzin": "PHEV-p",
            "(Plugin) Hybrid-Diesel": "PHEV-d",
            "Essence": "ICEV-p",
            "Gaz comprimé": "ICEV-g",
            "Electrique": "BEV",
            "Pile à combustible": "FCEV",
            "Hybride-essence": "HEV-p",
            "Hybride-diesel": "HEV-d",
            "Hybride-essence rechargeable": "PHEV-p",
            "Hybride-diesel rechargeable": "PHEV-d",
        }
        self.d_size_en = {
            "Minicompact": "Mini",
            "Subcompact": "Small",
            "Compact": "Lower medium",
            "Mid-size": "Medium",
            "Large": "Large",
            "SUV": "SUV",
            "Van": "Van",
        }
        self.d_size_fr = {
            "Mini-citadine": "Mini",
            "Citadine": "Small",
            "Berline compacte": "Lower medium",
            "Berline familiale": "Medium",
            "Grande routière": "Large",
            "SUV": "SUV",
            "Van": "Van",
        }
        self.d_size_it = {
            "Mini citycar": "Mini",
            "Citycar": "Small",
            "Berlina compatta": "Lower medium",
            "Berlina medio-grande": "Medium",
            "Berlina tre volumi": "Large",
            "SUV": "SUV",
            "Van": "Van",
        }
        self.d_size_de = {
            "Kleinstwagen": "Mini",
            "Kleinwagen": "Small",
            "Kompaktklasse": "Lower medium",
            "Mittelklasse": "Medium",
            "Oberklasse": "Large",
            "Geländewagen": "SUV",
            "Van": "Van",
        }
        self.d_size_all = {
            "Minicompact": "Mini",
            "Subcompact": "Small",
            "Compact": "Lower medium",
            "Mid-size": "Medium",
            "Large": "Large",
            "SUV": "SUV",
            "Van": "Van",
            "Mini-citadine": "Mini",
            "Citadine": "Small",
            "Berline compacte": "Lower medium",
            "Berline familiale": "Medium",
            "Grande routière": "Large",
            "Mini citycar": "Mini",
            "Citycar": "Small",
            "Berlina compatta": "Lower medium",
            "Berlina medio-grande": "Medium",
            "Berlina tre volumi": "Large",
            "Kleinstwagen": "Mini",
            "Kleinwagen": "Small",
            "Kompaktklasse": "Lower medium",
            "Mittelklasse": "Medium",
            "Oberklasse": "Large",
            "Geländewagen": "SUV",
        }

        self.d_rev_pt_en = {v: k for k, v, in self.d_pt_en.items()}
        self.d_rev_pt_fr = {v: k for k, v, in self.d_pt_fr.items()}
        self.d_rev_pt_it = {v: k for k, v, in self.d_pt_it.items()}
        self.d_rev_pt_de = {v: k for k, v, in self.d_pt_de.items()}

        self.d_rev_size_en = {v: k for k, v, in self.d_size_en.items()}
        self.d_rev_size_fr = {v: k for k, v, in self.d_size_fr.items()}
        self.d_rev_size_it = {v: k for k, v, in self.d_size_it.items()}
        self.d_rev_size_de = {v: k for k, v, in self.d_size_de.items()}

        self.excel = ""

    def load_map_file(self, lang):
        with open("data/car_to_class_map.csv", "r", encoding="ISO-8859-1") as f:
            data = [list(line) for line in csv.reader(f, delimiter=";")]

            if lang == "en":
                for d in data:
                    d[4] = self.d_rev_pt_en[d[4]]
                    d[5] = self.d_rev_size_en[d[5]]

            if lang == "fr":
                for d in data:
                    d[4] = self.d_rev_pt_fr[d[4]]
                    d[5] = self.d_rev_size_fr[d[5]]

            if lang == "de":
                for d in data:
                    d[4] = self.d_rev_pt_de[d[4]]
                    d[5] = self.d_rev_size_de[d[5]]

            if lang == "it":
                for d in data:
                    d[4] = self.d_rev_pt_it[d[4]]
                    d[5] = self.d_rev_size_it[d[5]]

        return data

    def load_params_file(self):
        with open("data/parameters definition.txt", "r") as f:
            data = [line for line in csv.reader(f, delimiter="\t")]
        return data

    def interpolate_array(self, years):
        return self.arr.interp(year=years, kwargs={"fill_value": "extrapolate"})

    def get_dc(self, dc):
        return get_standard_driving_cycle(dc)

    def create_config_array(self, dict_params, array, mix, energy_storage):

        arr = []
        year = [int(y) for y in dict_params[("Functional unit",)]["year"]]
        driving_cycle = dict_params[("Driving cycle",)]
        country = dict_params[("Background",)]["country"]
        passengers = dict_params[("Foreground",)][
            ("Glider", "all", "all", "average passengers", "none")
        ][(year[0], "loc")]
        cargo_mass = dict_params[("Foreground",)][
            ("Glider", "all", "all", "cargo mass", "none")
        ][(year[0], "loc")]
        lifetime = dict_params[("Foreground",)][
            ("Driving", "all", "all", "lifetime kilometers", "none")
        ][(year[0], "loc")]
        km_per_year = dict_params[("Foreground",)][
            ("Driving", "all", "all", "kilometers per year", "none")
        ][(year[0], "loc")]

        for pt in array.coords["powertrain"].values:
            for s in array.coords["size"].values:
                for y, year in enumerate(array.coords["year"].values.astype(int)):

                    electricity_mix = mix[y].tolist()

                    params = [
                        pt,
                        s,
                        int(year),
                        lifetime,
                        km_per_year,
                        passengers,
                        cargo_mass,
                        driving_cycle,
                        country,
                        electricity_mix,
                    ]
                    other_params = (
                        array.sel(
                            powertrain=pt,
                            size=s,
                            year=year,
                            value=0,
                            parameter=[
                                "TtW energy",
                                "driving mass",
                                "combustion power",
                                "electric power",
                                "range",
                                "engine efficiency",
                                "drivetrain efficiency",
                                "TtW efficiency",
                                "battery discharge efficiency",
                                "energy battery mass",
                                "battery cell energy density",
                                "electric energy stored",
                                "battery lifetime kilometers",
                            ],
                        )
                        .values.astype(float)
                        .tolist()
                    )
                    params.extend(other_params)

                    if pt in ("BEV"):
                        battery_chem = dict_params[("Background",)]["energy storage"][
                            "electric"
                        ]["type"]
                        battery_origin = dict_params[("Background",)]["energy storage"][
                            "electric"
                        ]["origin"]
                        (
                            primary_fuel_type,
                            primary_fuel_share,
                            secondary_fuel_type,
                            secondary_fuel_share,
                        ) = ["", "", "", ""]
                    else:
                        battery_chem, battery_origin = ["", ""]

                    if pt in ("ICEV-p", "PHEV-p", "HEV-p"):
                        if "fuel blend" in dict_params[("Background",)]:
                            if "petrol" in dict_params[("Background",)]["fuel blend"]:
                                primary_fuel_type = dict_params[("Background",)][
                                    "fuel blend"
                                ]["petrol"]["primary fuel"]["type"]
                                primary_fuel_share = dict_params[("Background",)][
                                    "fuel blend"
                                ]["petrol"]["primary fuel"]["share"][y]
                                secondary_fuel_type = dict_params[("Background",)][
                                    "fuel blend"
                                ]["petrol"]["secondary fuel"]["type"]
                                secondary_fuel_share = dict_params[("Background",)][
                                    "fuel blend"
                                ]["petrol"]["secondary fuel"]["share"][y]
                            else:
                                if country in self.biogasoline.country.values:
                                    share_biogasoline = np.squeeze(np.clip(
                                        self.biogasoline.sel(
                                            country=country
                                        )
                                            .interp(year=year, kwargs={"fill_value": "extrapolate"})
                                            .values
                                        , 0, 1)).tolist()
                                else:
                                    share_biogasoline = 0
                                (
                                    primary_fuel_type,
                                    primary_fuel_share,
                                    secondary_fuel_type,
                                    secondary_fuel_share,
                                ) = [
                                    "petrol",
                                    1 - share_biogasoline,
                                    "bioethanol - wheat straw",
                                    share_biogasoline,
                                ]

                        else:
                            if country in self.biogasoline.country.values:
                                share_biogasoline = np.squeeze(np.clip(
                                    self.biogasoline.sel(
                                        country=country
                                    )
                                        .interp(year=year, kwargs={"fill_value": "extrapolate"})
                                        .values
                                    , 0, 1)).tolist()
                            else:
                                share_biogasoline = 0
                            (
                                primary_fuel_type,
                                primary_fuel_share,
                                secondary_fuel_type,
                                secondary_fuel_share,
                            ) = [
                                "petrol",
                                1 - share_biogasoline,
                                "bioethanol - wheat straw",
                                share_biogasoline,
                            ]

                    if pt in ("ICEV-d", "PHEV-d", "HEV-d"):
                        if "fuel blend" in dict_params[("Background",)]:
                            if "diesel" in dict_params[("Background",)]["fuel blend"]:
                                primary_fuel_type = dict_params[("Background",)][
                                    "fuel blend"
                                ]["diesel"]["primary fuel"]["type"]
                                primary_fuel_share = dict_params[("Background",)][
                                    "fuel blend"
                                ]["diesel"]["primary fuel"]["share"][y]
                                secondary_fuel_type = dict_params[("Background",)][
                                    "fuel blend"
                                ]["diesel"]["secondary fuel"]["type"]
                                secondary_fuel_share = dict_params[("Background",)][
                                    "fuel blend"
                                ]["diesel"]["secondary fuel"]["share"][y]
                            else:
                                if country in self.biodiesel.country.values:
                                    share_biodiesel = np.squeeze(np.clip(
                                        self.biodiesel.sel(
                                            country=country
                                        )
                                            .interp(year=year, kwargs={"fill_value": "extrapolate"})
                                            .values
                                        , 0, 1)).tolist()
                                else:
                                    share_biodiesel = 0
                                (
                                    primary_fuel_type,
                                    primary_fuel_share,
                                    secondary_fuel_type,
                                    secondary_fuel_share,
                                ) = [
                                    "diesel",
                                    1 - share_biodiesel,
                                    "biodiesel - cooking oil",
                                    share_biodiesel,
                                ]

                        else:
                            if country in self.biodiesel.country.values:
                                share_biodiesel = np.squeeze(np.clip(
                                    self.biodiesel.sel(
                                        country=country
                                    )
                                        .interp(year=year, kwargs={"fill_value": "extrapolate"})
                                        .values
                                    , 0, 1)).tolist()
                            else:
                                share_biodiesel = 0
                            (
                                primary_fuel_type,
                                primary_fuel_share,
                                secondary_fuel_type,
                                secondary_fuel_share,
                            ) = [
                                "diesel",
                                1 - share_biodiesel,
                                "biodiesel - cooking oil",
                                share_biodiesel,
                            ]

                    if pt in ("ICEV-g"):
                        if "fuel blend" in dict_params[("Background",)]:
                            if "cng" in dict_params[("Background",)]["fuel blend"]:
                                primary_fuel_type = dict_params[("Background",)][
                                    "fuel blend"
                                ]["cng"]["primary fuel"]["type"]
                                primary_fuel_share = dict_params[("Background",)][
                                    "fuel blend"
                                ]["cng"]["primary fuel"]["share"][y]
                                secondary_fuel_type = dict_params[("Background",)][
                                    "fuel blend"
                                ]["cng"]["secondary fuel"]["type"]
                                secondary_fuel_share = dict_params[("Background",)][
                                    "fuel blend"
                                ]["cng"]["secondary fuel"]["share"][y]
                            else:
                                if country in self.biomethane.country.values:
                                    share_biomethane = np.squeeze(np.clip(
                                        self.biomethane.sel(
                                            country=country
                                        )
                                            .interp(year=year, kwargs={"fill_value": "extrapolate"})
                                            .values
                                        , 0, 1)).tolist()
                                else:
                                    share_biomethane = 0
                                (
                                    primary_fuel_type,
                                    primary_fuel_share,
                                    secondary_fuel_type,
                                    secondary_fuel_share,
                                ) = [
                                    "cng",
                                    1 - share_biomethane,
                                    "biogas - sewage sludge",
                                    share_biomethane,
                                ]

                        else:
                            if country in self.biomethane.country.values:
                                share_biomethane = np.squeeze(np.clip(
                                    self.biomethane.sel(
                                        country=country
                                    )
                                        .interp(year=year, kwargs={"fill_value": "extrapolate"})
                                        .values
                                    , 0, 1)).tolist()
                            else:
                                share_biomethane = 0
                            (
                                primary_fuel_type,
                                primary_fuel_share,
                                secondary_fuel_type,
                                secondary_fuel_share,
                            ) = [
                                "cng",
                                1 - share_biomethane,
                                "biogas - sewage sludge",
                                share_biomethane,
                            ]

                    if pt in ("FCEV"):
                        if "fuel blend" in dict_params[("Background",)]:
                            if "hydrogen" in dict_params[("Background",)]["fuel blend"]:
                                primary_fuel_type = dict_params[("Background",)][
                                    "fuel blend"
                                ]["hydrogen"]["primary fuel"]["type"]
                                primary_fuel_share = dict_params[("Background",)][
                                    "fuel blend"
                                ]["hydrogen"]["primary fuel"]["share"][y]
                                secondary_fuel_type = dict_params[("Background",)][
                                    "fuel blend"
                                ]["hydrogen"]["secondary fuel"]["type"]
                                secondary_fuel_share = dict_params[("Background",)][
                                    "fuel blend"
                                ]["hydrogen"]["secondary fuel"]["share"][y]
                            else:
                                (
                                    primary_fuel_type,
                                    primary_fuel_share,
                                    secondary_fuel_type,
                                    secondary_fuel_share,
                                ) = ["electrolysis", 1, "", ""]
                        else:
                            (
                                primary_fuel_type,
                                primary_fuel_share,
                                secondary_fuel_type,
                                secondary_fuel_share,
                            ) = ["electrolysis", 1, "", ""]

                    params.extend(
                        [
                            battery_chem,
                            battery_origin,
                            primary_fuel_type,
                            primary_fuel_share,
                            secondary_fuel_type,
                            secondary_fuel_share,
                        ]
                    )

                    arr.append(params)

        return arr

    def process_results(self, d, lang, job_id):
        """ Calculate LCIA and store results in an array of arrays """

        # Update task progress to db
        task = Task.query.filter_by(id=job_id).first()
        task.progress = 50
        db.session.commit()

        scope = {
            "powertrain": d[("Functional unit",)]["powertrain"],
            "size": d[("Functional unit",)]["size"],
        }

        self.dcts, self.arr = fill_xarray_from_input_parameters(self.cip, scope=scope)
        arr = self.interpolate_array(d[("Functional unit",)]["year"])
        modify_xarray_from_custom_parameters(d[("Foreground",)], arr)

        # remove hybridization for vehicles before 2030
        pwt = list({"ICEV-p", "ICEV-d", "ICEV-g"}.intersection(set(scope["powertrain"])))
        years_before_2030 = [y for y in arr["year"].values if y < 2030]

        if pwt and years_before_2030:
            arr.loc[dict(
                powertrain=pwt, year=years_before_2030, parameter="combustion power share"
            )] = 1

        cm = CarModel(arr, cycle=d[("Driving cycle",)])

        # adjust the electricity density of the battery cells
        for p in d[('Foreground',)]:
            if p[3] == "battery cell energy density":
                for y in d[("Foreground",)][p]:
                    cm.array.loc[
                        dict(
                            parameter="battery cell energy density",
                            year=y[0]
                        )
                    ]= d[("Foreground",)][p][y]

        if "electric utility factor" in d[("Background",)]:
            uf = list(d[("Background",)]["electric utility factor"].values())
            cm.set_all(electric_utility_factor=uf)
        else:
            cm.set_all()

        pt = cm.array.powertrain.values
        s = d[("Functional unit",)]["size"]
        y = d[("Functional unit",)]["year"]
        a = [pt] + [s] + [y]
        l = list(itertools.product(*a))
        l = [i[0] + " - " + i[1] + " - " + str(i[2]) for i in l]

        cumsum = (
            cm.energy.sel(
                powertrain=pt,
                size=s,
                year=y,
                value=0,
                parameter=["motive energy", "auxiliary energy", "recuperated energy"],
            )
            .cumsum(dim="second")
            .sum(dim="parameter")
            .transpose("powertrain", "size", "year", "second")
            .values.reshape(len(l), -1).astype("float64")
        )

        # Format the data so that it can be consumed directly
        # by nvd3.js
        TtW_energy = []
        for i, vehicle in enumerate(l):
            TtW_energy.append(
                {
                    "key": vehicle,
                    "values": list(
                        map(lambda e: {"x": e[0], "y": e[1]}, enumerate(cumsum[i]))
                    ),
                }
            )

        # Functional unit
        fu_unit = d[("Functional unit",)]["fu"]["unit"]
        fu_qty = float(d[("Functional unit",)]["fu"]["quantity"])

        if fu_unit == "vkm":
            load_factor = 1
        else:
            load_factor = cm["average passengers"].mean().values

        # Update task progress to db
        task = Task.query.filter_by(id=job_id).first()
        task.progress = 60
        db.session.commit()

        scope = {"powertrain": pt, "size": s, "year": y}
        total_cost = cm.calculate_cost_impacts(scope=scope).transpose(
            "size", "powertrain", "year", "value", "cost_type"
        ).astype("float64")

        cost_benchmark = total_cost.sel(cost_type="total", value=0).values.reshape(
            len(l)
        )

        cost_types = [c for c in total_cost.cost_type.values if c != "total"]

        arr_benchmark = list(
            map(
                lambda x: [
                    "cost",
                    x[0].split(" - ")[0],
                    x[0].split(" - ")[1],
                    x[0].split(" - ")[2],
                    1 / x[1],
                ],
                zip(l, cost_benchmark),
            )
        )

        l_scatter = [x.replace(" - ", ", ") for x in l]

        dict_scatter = {
            x[0]: [x[1]]
            for x in zip(l_scatter, cost_benchmark / load_factor * fu_qty)
        }

        detailed_cost = (
            total_cost.sel(value=0, cost_type=cost_types).values.reshape(
                len(l), len(cost_types)
            )
            / load_factor
            * fu_qty
        )

        a = [pt] + [s] + [y]
        l_cost = list(itertools.product(*a))

        list_res_costs = list(
            map(
                lambda x: [
                    [
                        "ownership cost",
                        x[0][1],
                        x[0][0],
                        x[0][2],
                        cost_types[y],
                        z,
                        np.sum(x[1]),
                    ]
                    for y, z in enumerate(x[1])
                ],
                zip(l_cost, detailed_cost),
            )
        )

        list_res_costs = list(itertools.chain.from_iterable(list_res_costs))

        self.ic = InventoryCalculation(
            cm.array,
            scope=d[("Functional unit",)]["fu"],
            background_configuration=d[("Background",)],
        )
        # Update task progress to db
        task = Task.query.filter_by(id=job_id).first()
        task.progress = 70
        db.session.commit()

        results = (
            self.ic.calculate_impacts()
            .sel(value=0)
            .transpose("impact_category", "size", "powertrain", "year", "impact")
        ).astype("float64")

        lifetime = int(cm.array.sel(parameter="lifetime kilometers").mean().values)

        # Update task progress to db
        task = Task.query.filter_by(id=job_id).first()
        task.progress = 80
        db.session.commit()

        impact = results.coords["impact"].values.tolist()
        impact_category = results.coords["impact_category"].values

        arr_benchmark.extend(
            list(
                map(
                    lambda x: [
                        "climate change",
                        x[0].split(" - ")[0],
                        x[0].split(" - ")[1],
                        x[0].split(" - ")[2],
                        1 / x[1],
                    ],
                    zip(
                        l,
                        results.sel(impact_category="climate change")
                        .sum(dim="impact")
                        .values.reshape(len(l)),
                    ),
                )
            )
        )

        arr_benchmark.extend(
            list(
                map(
                    lambda x: [
                        "fossil depletion",
                        x[0].split(" - ")[0],
                        x[0].split(" - ")[1],
                        x[0].split(" - ")[2],
                        1 / x[1] * 0.755,  # 0.755 kg/L gasoline
                    ],
                    zip(
                        l,
                        results.sel(impact_category="fossil depletion")
                        .sum(dim="impact")
                        .values.reshape(len(l)),
                    ),
                )
            )
        )

        for x in zip(
                l_scatter,
                results.sel(impact_category="climate change")
                        .sum(dim="impact")
                        .values.reshape(len(l))
                / load_factor
                * fu_qty,
            ):
            existing_list = dict_scatter[x[0]]
            existing_list.append(x[1])
            dict_scatter[x[0]] = existing_list

        a_wo_impact = [impact_category] + [s] + [pt] + [y]
        l_impacts_wo_impact = list(itertools.product(*a_wo_impact))

        list_res = list(
            map(
                lambda x: [
                    [x[0][0], x[0][1], x[0][2], x[0][3], impact[y], z, np.sum(x[1])]
                    for y, z in enumerate(x[1])
                ],
                zip(
                    l_impacts_wo_impact,
                    (
                        results.values.reshape(
                            len(
                                l_impacts_wo_impact
                            ),
                            len(impact)
                        )
                        / load_factor
                        * fu_qty
                    ),
                ),
            )
        )

        list_res = list(itertools.chain.from_iterable(list_res))

        list_res_acc = list(
            map(
                lambda x: [
                    x[0][0],
                    x[0][1],
                    x[0][2],
                    x[0][3],
                    np.sum(x[1][4:-1]) * lifetime,
                    np.sum(x[1][[0, 1, 2, 3, -1]]),
                    lifetime,
                ],
                zip(
                    l_impacts_wo_impact,
                    (
                        results.values.reshape(
                            len(
                                l_impacts_wo_impact
                            ),
                            len(impact),
                        )
                        / load_factor
                        * fu_qty
                    ),
                ),
            )
        )

        # Update task progress to db
        task = Task.query.filter_by(id=job_id).first()
        task.progress = 90
        db.session.commit()

        self.ic = InventoryCalculation(
            cm.array,
            scope=d[("Functional unit",)]["fu"],
            background_configuration=d[("Background",)],
            method="ilcd",
        )
        results = self.ic.calculate_impacts().astype("float64")

        nf = np.array(
            [
                1.18e4,
                4.75e-4,
                3.85e-5,
                6.36e-2,
                8.4e3,
                8.4e3,
                8.4e3,
                8.4e3,
                5.55e1,
                7.34e-1,
                2.83e1,
                1.77e2,
                4.22e3,
                2.34e-2,
                4.06e1,
                7.18e-4,
                1.18e4,
                6.53e4,
                1.4e6,
            ]
        )

        nf_impact = (results / nf[:, None, None, None, None, None]).sum(dim="impact")
        impact_category = nf_impact.coords["impact_category"].values

        a = [impact_category] + [s] + [pt] + [y]
        l_norm = list(itertools.product(*a))

        list_normalized_results = list(
            map(
                lambda x: [x[0][0], x[0][1], x[0][2], x[0][3], x[1]],
                zip(
                    l_norm,
                    (nf_impact.values.reshape(len(l_norm)) / load_factor * fu_qty),
                ),
            )
        )

        # Update task progress to db
        task = Task.query.filter_by(id=job_id).first()
        task.progress = 85
        db.session.commit()

        # reformat A for export
        self.ic.inputs = self.ic.get_dict_input()
        #self.ic.bs = BackgroundSystemModel()
        self.ic.country = self.ic.get_country_of_use()
        self.ic.add_additional_activities()
        self.ic.rev_inputs = self.ic.get_rev_dict_input()
        self.ic.A = self.ic.get_A_matrix()

        # add vehicles datasets
        self.ic.add_additional_activities_for_export()

        # Update dictionary
        self.ic.rev_inputs = self.ic.get_rev_dict_input()

        # resize A matrix
        self.ic.A = self.ic.get_A_matrix()

        # Create electricity and fuel market datasets
        self.ic.create_electricity_market_for_fuel_prep()

        # Create fuel markets
        self.ic.fuel_blends = {}
        self.ic.fuel_dictionary = self.ic.create_fuel_dictionary()
        self.ic.define_fuel_blends()
        self.ic.set_actual_range()

        # Create electricity market dataset for battery production
        self.ic.create_electricity_market_for_battery_production()
        self.ic.set_inputs_in_A_matrix_for_export(self.ic.array.values)

        self.export = ExportInventory(self.ic.A, self.ic.rev_inputs)

        # Update task progress to db
        task = Task.query.filter_by(id=job_id).first()
        task.progress = 100
        db.session.commit()

        list_res.extend(list_res_costs)

        return (
            json.dumps(
                [
                    lang,
                    list_res,
                    arr_benchmark,
                    TtW_energy,
                    dict_scatter,
                    list_res_acc,
                    self.create_config_array(d, cm.array, self.ic.mix, cm.energy_storage),
                    d[("Background",)]["country"],
                    d[("Functional unit",)]["fu"]["quantity"],
                    d[("Functional unit",)]["fu"]["unit"],
                    list_normalized_results,
                ]
            ),
            self.export,
        )

    def format_dictionary(self, raw_dict, lang, job_id):
        """ Format the dictionary sent by the user so that it can be understood by `carculator` """

        d_sliders = {
            "mileage-slider": "kilometers per year",
            "lifetime-slider": "lifetime kilometers",
            "passenger-slider": "average passengers",
            "cargo-slider": "cargo mass",
        }
        new_dict = {}

        new_dict[("Functional unit",)] = {
            "powertrain": [x for x in raw_dict["type"]],
            "year": [int(x) for x in raw_dict["year"]],
            "size": [s for s in raw_dict["size"]],
            "fu": raw_dict["fu"],
        }

        f_d = {}
        new_dict[("Driving cycle",)] = raw_dict["driving_cycle"]
        new_dict[("Background",)] = {
            k: v
            for k, v in raw_dict["background params"].items()
            if k not in ("energy storage", "efficiency")
        }

        # Ensure that the electricity mix is present
        if "custom electricity mix" not in new_dict[("Background",)]:
            years = new_dict[("Functional unit",)]["year"]
            country = new_dict[("Background",)]["country"]
            try:
                response = (
                    self.electricity_mix.loc[
                        dict(
                            country=country,
                            variable=[
                                "Hydro",
                                "Nuclear",
                                "Gas",
                                "Solar",
                                "Wind",
                                "Biomass",
                                "Coal",
                                "Oil",
                                "Geothermal",
                                "Waste",
                            ],
                        )
                    ]
                    .interp(year=years)
                    .values
                )
            except KeyError:
                response = (
                    self.electricity_mix.loc[
                        dict(
                            country="RER",
                            variable=[
                                "Hydro",
                                "Nuclear",
                                "Gas",
                                "Solar",
                                "Wind",
                                "Biomass",
                                "Coal",
                                "Oil",
                                "Geothermal",
                                "Waste",
                            ],
                        )
                    ]
                    .interp(year=years)
                    .values
                )

            new_dict[("Background",)]["custom electricity mix"] = response

        # Ensure that the electricity mix split equals 1
        for el in new_dict[("Background",)]["custom electricity mix"]:
            el /= np.sum(np.array(el))

        if "energy storage" in raw_dict["background params"]:
            if "electric" in raw_dict["background params"]["energy storage"]:
                if len(raw_dict["background params"]["energy storage"]["electric"]) > 0:
                    energy_storage = raw_dict["background params"]["energy storage"][
                        "electric"
                    ]
                    new_dict[("Background",)]["energy storage"] = {"electric": {}}

                    for e in energy_storage:
                        for p in energy_storage[e]:
                            if p in ("type", "origin"):
                                new_dict[("Background",)]["energy storage"]["electric"][
                                    p
                                ] = energy_storage[e][p]

        for k, v in raw_dict["foreground params"].items():
            if k in d_sliders:
                name = d_sliders[k]
                cat = self.d_categories[name]
                powertrain = "all"
                size = "all"

                if isinstance(v, str):
                    val = [float(v.replace(" ", ""))] * len(
                        new_dict[("Functional unit",)]["year"]
                    )
                else:
                    val = [float(v)] * len(
                        new_dict[("Functional unit",)]["year"]
                    )
            else:

                k = tuple(k.split(","))
                name = k[0]
                cat = self.d_categories[name]
                powertrain = k[1]
                size = k[2]
                val = [float(n) for n in v] * len(
                    new_dict[("Functional unit",)]["year"]
                )

            d_val = {
                (k, "loc"): v
                for k, v in list(zip(new_dict[("Functional unit",)]["year"], val))
            }

            f_d[(cat, powertrain, size, name, "none")] = d_val

        if "energy storage" in raw_dict["background params"]:
            if "electric" in raw_dict["background params"]["energy storage"]:
                energy_storage = raw_dict["background params"]["energy storage"][
                    "electric"
                ]

                for e in energy_storage:
                    size = e
                    for p in energy_storage[e]:
                        if p not in ("type", "origin"):

                            cat = self.d_categories[name]
                            powertrain = "BEV"
                            val = energy_storage[e][p]

                            d_val = {
                                (k, "loc"): v
                                for k, v in list(
                                    zip(new_dict[("Functional unit",)]["year"], val)
                                )
                            }

                            f_d[(cat, powertrain, size, p, "none")] = d_val

        if "efficiency" in raw_dict["background params"]:
            efficiency = raw_dict["background params"]["efficiency"]
            efficiency = {
                k: v for k, v in efficiency.items() if len(list(v.keys())) > 0
            }

            for eff in efficiency:
                powertrain = eff
                for s in efficiency[eff]:
                    size = s
                    for c in efficiency[eff][s]:
                        name = c
                        cat = self.d_categories[name]

                        if np.sum(efficiency[eff][s][c]) != 0:
                            val = efficiency[eff][s][c]
                            d_val = {
                                (k, "loc"): v
                                for k, v in list(
                                    zip(new_dict[("Functional unit",)]["year"], val)
                                )
                            }
                            f_d[(cat, powertrain, size, name, "none")] = d_val

        new_dict[("Foreground",)] = f_d

        return new_dict


'''
 Excited States software: qFit 3.0

 Contributors: Saulo H. P. de Oliveira, Gydo van Zundert, and Henry van den Bedem.
 Contact: vdbedem@stanford.edu

 Copyright (C) 2009-2019 Stanford University
 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:

 This entire text, including 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, CONTRIBUTORS 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 numpy as np
import copy
import math
from .base_structure import _BaseStructure
from .math import *
from .rotamers import ROTAMERS
import time

_PROTEIN_BACKBONE_ATOMS = ['N', 'CA', 'C']
_NUCLEOTIDE_BACKBONE_ATOMS = ['P', "O5'", "C5'", "C4'", "C3'", "O3"]
_SOLVENTS = ['HOH']

def residue_type(residue):
    """Check residue type.

    The following types are recognized:
        rotamer-residue
            Amino acid for which rotamer information is available
        aa-residue
            Amino acid residue
        residue
            DNA / RNA residue
        solvent
            HOH
        ligand
            Other
    """

    # RNA and DNA backbone atoms
    if residue.resn[0] in ROTAMERS:
        return 'rotamer-residue'
    bb_atoms = _PROTEIN_BACKBONE_ATOMS
    if residue.select('name', bb_atoms).size == 3:
        return 'aa-residue'
    xna_atoms = _NUCLEOTIDE_BACKBONE_ATOMS
    if residue.select('name', xna_atoms).size == 6:
        return 'residue'
    if residue.resn[0] in _SOLVENTS:
        return 'solvent'
    return 'ligand'

class _BaseResidue(_BaseStructure):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.id = (kwargs['resi'], kwargs['icode'])
        self.type = kwargs['type']

    def __repr__(self):
        resi, icode = self.id
        string = "Residue: {}".format(resi)
        if self.id[1] != "":
            string += ":{}".format(icode)
        return string

class _Residue(_BaseResidue):
    pass

class _RotamerResidue(_BaseResidue):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        resname = self.resn[0]
        self._rotamers = ROTAMERS[resname]
        self.nchi = self._rotamers['nchi']
        self.nrotamers = len(self._rotamers['rotamers'])
        self.rotamers = self._rotamers['rotamers']
        self._init_clash_detection()
        self.i=0

    def _init_clash_detection(self, scaling_factor=0.75, covalent_bonds=None):
        # Setup the condensed distance based arrays for clash detection and fill them
        self._ndistances = self.natoms * (self.natoms - 1) // 2
        self._clash_mask = np.ones(self._ndistances, bool)
        self._clash_radius2 = np.zeros(self._ndistances, float)
        radii = self.covalent_radius
        bonds = self._rotamers['bonds']
        if covalent_bonds:
            for bond in covalent_bonds:
                bonds.append(bond)
        offset = self.natoms * (self.natoms - 1) // 2
        for i in range(self.natoms - 1):
            #starting_index = int(offset - sp_comb(self.natoms - i, 2)) - i - 1
            natoms_left = self.natoms - i
            starting_index = offset - natoms_left * (natoms_left - 1) // 2 - i - 1
            name1 = self.name[i]
            covrad1 = radii[i]
            for j in range(i + 1, self.natoms):
                bond1 = [name1, self.name[j]]
                bond2 = bond1[::-1]
                covrad2 = radii[j]
                index = starting_index + j
                self._clash_radius2[index] = covrad1 + covrad2 + 0.4
                if bond1 in bonds or bond2 in bonds:
                    self._clash_mask[index] = False

        self._clash_radius2 *= self._clash_radius2
        self._clash_radius2 *= scaling_factor
        self._clashing = np.zeros(self._ndistances, bool)
        self._dist2_matrix = np.empty(self._ndistances, float)

        # All atoms are active from the start
        self.active = np.ones(self.natoms, bool)
        self._active_mask = np.ones(self._ndistances, bool)

    def update_clash_mask(self):
        offset = self.natoms * (self.natoms - 1) // 2 - 1
        for i, active in enumerate(self.active[:-1]):
            natoms_left = self.natoms - i
            starting_index = offset - natoms_left * (natoms_left - 1) // 2 - i
            end = starting_index + self.natoms - i + 1
            self._active_mask[starting_index: end] = active

    def clashes(self):
        """Checks if there are any internal clashes.
        Deactivated atoms are not taken into account.
        """

        dm = self._dist2_matrix
        coor = self.coor
        dot = np.dot
        k = 0
        for i in range(self.natoms - 1):
            u = coor[i]
            for j in range(i + 1, self.natoms):
                u_v = u - coor[j]
                dm[k] = dot(u_v, u_v)
                k += 1

        np.less_equal(dm, self._clash_radius2 , self._clashing)
        self._clashing &= self._clash_mask
        self._clashing &= self._active_mask
        nclashes = self._clashing.sum()
        return nclashes

    def get_chi(self, chi_index):
        atoms = self._rotamers['chi'][chi_index]
        selection = self.select('name', atoms)
        ordered_sel = []
        for atom in atoms:
            for sel in selection:
                if atom == self._name[sel]:
                    ordered_sel.append(sel)
                    break
        coor = self._coor[ordered_sel]
        angle = dihedral_angle(coor)
        return angle

    def set_chi(self, chi_index, value, covalent=None, length=None):
        atoms = self._rotamers['chi'][chi_index]
        selection = self.select('name', atoms)
        coor = self._coor[selection]
        origin = coor[1].copy()
        coor -= origin
        zaxis = coor[2]
        zaxis /= np.linalg.norm(zaxis)
        yaxis = coor[0] - np.inner(coor[0], zaxis) * zaxis
        yaxis /= np.linalg.norm(yaxis)
        xaxis = np.cross(yaxis, zaxis)
        backward = np.asmatrix(np.zeros((3, 3), float))
        backward[0] = xaxis
        backward[1] = yaxis
        backward[2] = zaxis
        forward = backward.T

        atoms_to_rotate = self._rotamers['chi-rotate'][chi_index]
        selection = self.select('name', atoms_to_rotate)
        if covalent in atoms_to_rotate:
            # If we are rotating the atom that is covalently bonded
            # to the ligand, we should also rotate the ligand.
            atoms_to_rotate2 = self.name[length:]
            selection2 = self.select('name', atoms_to_rotate2)
            tmp = list(selection)
            tmp += list(selection2)
            selection = np.array(tmp, dtype=int)
        coor_to_rotate = np.dot(self._coor[selection] - origin, backward.T)
        rotation = Rz(np.deg2rad(value - self.get_chi(chi_index)))

        R = forward * rotation
        self._coor[selection] = coor_to_rotate.dot(R.T) + origin

    def print_residue(self):
        for atom, coor, element, b, q in zip(
          self.name, self.coor, self.e, self.b, self.q):
            print("{} {} {} {} {}".format(atom, coor, element, b, q))

    def complete_residue(self):
        self.visited = []
        if residue_type(self) != "rotamer-residue":
            msg = ("Cannot complete non-aminoacid residue. Please, "
                   "complete the missing atoms of the residue for qFiting!")
            raise RuntimeError(msg)

        for atom, position in zip(self._rotamers['atoms'],
                                  self._rotamers['positions']):
            # Found a missing atom!:
            if atom not in self.name:
                self.complete_residue_recursive(atom)

    def complete_residue_recursive(self, atom):
        if atom in ['N', 'C', 'CA', 'O']:
            msg = ("Cannot complete missing backbone atoms. Please, "
                   "complete the missing backbone atoms of the residue"
                   " for qFiting!")
            raise RuntimeError(msg)
        ref_atom = self._rotamers['connectivity'][atom][0]
        if ref_atom not in self.name:
            self.complete_residue_recursive(ref_atom)
        idx = np.argwhere(self.name == ref_atom)[0]
        ref_coor = self.coor[idx]
        bond_length, bond_length_sd = (
            self._rotamers['bond_dist'][ref_atom][atom])
        # Identify a suitable atom for the bond angle:
        for angle in self._rotamers['bond_angle']:
            if angle[0][1] == ref_atom and angle[0][2] == atom:
                if angle[0][0][0] == "H":
                    continue
                bond_angle_atom = angle[0][0]
                bond_angle, bond_angle_sd = angle[1]
                if bond_angle_atom not in self.name:
                    self.complete_residue_recursive(bond_angle_atom)
                bond_angle_coor = self.coor[np.argwhere(
                  self.name == bond_angle_atom)[0]]
                dihedral_atom = None
                # If the atom's position is dependent on a rotamer,
                # identify the fourth dihedral angle atom:
                for i, chi in self._rotamers['chi'].items():
                    if (chi[1] == bond_angle_atom and
                       chi[2] == ref_atom and
                       chi[3] == atom):
                        dihedral_atom = chi[0]
                        dihed_angle = self._rotamers['rotamers'][0][i-1]

                # If the atom's position is not dependent on a rotamer,
                # identify the fourth dihedral angle atom:
                if dihedral_atom is None:
                    for dihedral in self._rotamers['dihedral']:
                        if (dihedral[0][1] == bond_angle_atom and
                           dihedral[0][2] == ref_atom and
                           dihedral[0][3] == atom):
                            dihedral_atom = dihedral[0][0]
                            if dihedral[1][0] in self._rotamers['atoms']:
                                other_dihedral_atom = dihedral[1][0]
                                if dihedral_atom not in self.name:
                                    self.complete_residue_recursive(
                                     dihedral_atom)
                                dihedral_atom_coor = self.coor[np.argwhere(
                                 self.name == dihedral_atom)[0]]
                                if other_dihedral_atom not in self.name:
                                    self.complete_residue_recursive(
                                     other_dihedral_atom)
                                other_dihedral_atom_coor = self.coor[
                                 np.argwhere(
                                  self.name == other_dihedral_atom)[0]]
                                try:
                                    dihed_angle = dihedral[1][1]
                                    dihed_angle += dihedral_angle([
                                                    dihedral_atom_coor[0],
                                                    bond_angle_coor[0],
                                                    ref_coor[0],
                                                    other_dihedral_atom_coor[0]
                                                    ])
                                except:
                                    dihed_angle = 180
                                    dihed_angle += dihedral_angle([
                                                    dihedral_atom_coor[0],
                                                    bond_angle_coor[0],
                                                    ref_coor[0],
                                                    other_dihedral_atom_coor[0]
                                                    ])
                            else:
                                dihed_angle = dihedral[1][0]
                            break
                if dihedral_atom is not None:
                    if dihedral_atom not in self.name:
                        self.complete_residue_recursive(dihedral_atom)
                    dihedral_atom_coor = self.coor[np.argwhere(
                                            self.name == dihedral_atom)[0]]
                    break
        new_coor = self.calc_coordinates(dihedral_atom_coor[0],
                                         bond_angle_coor[0],
                                         ref_coor[0],
                                         bond_length,
                                         bond_angle,
                                         dihed_angle)
        new_coor = [round(x, 3) for x in new_coor]
        self.add_atom(atom, atom[0], new_coor)

    def calc_coordinates(self,u, v, Origin, L, bond_angle, dihedral):
        a=u-Origin
        b=v-Origin
        # Find the normal vector that defines the plane between a and b
        n = np.cross(a, b)
        AX,AY,AZ = list(a) ## a vector
        BX,BY,BZ = list(b) ## b vector
        A ,B ,C  = list(n) ## n vector
        # CALCULATE COORDINATES OF ATOM IF POSITIONED COPLANAR TO BOND ATOMS AT LENGTH L AND ANGLE BOND_ANGLE
        F = np.linalg.norm(b)  * np.cos(np.deg2rad(bond_angle))
        denom= (B*B)*(BX*BX+BZ*BZ)+ (A*A)*(BY*BY+BZ*BZ) + (BX*BX+BY*BY)*(C*C) - (2*A*BX*BZ*C) - (2*B*BY)*(A*BX+BZ*C)
        const= L * np.abs(B*BZ-BY*C) * np.sqrt(-F * F * np.inner(n,n)+denom)
        X= (  L*F*((B*B*BX)-(A*B*BY)+C*(-A*BZ+BX*C)) + const )/denom
        if((B==0 or BZ==0) and (BY==0 or C==0)):
            const1=math.sqrt( C*C*(-A*A*X*X+(B*B+C*C)*(L-X)*(L+X)))
            Y= ((-A*B*X)+const1)/(B*B+C*C)
            Z= -(A*C*C*X+B*const1)/(C*(B*B+C*C))
        else:
            Y= ((A*A*BY*F*L)*(B*BZ-BY*C)+ C*( -F*L*math.pow(B*BZ-BY*C,2) + BX*const) - A*( B*B*BX*BZ*F*L- B*BX*BY*F*L*C + BZ*const)) / ((B*BZ-BY*C)*denom)
            Z= ((A*A*BZ*F*L)*(B*BZ-BY*C) + (B*F*L)*math.pow(B*BZ-BY*C,2) + (A*BX*F*L*C)*(-B*BZ+BY*C) - B*BX*const + A*BY*const) / ((B*BZ-BY*C)*denom)
        # Translate the new vector to the correct coordinate
        D=np.array([X, Y, Z]) + Origin
        # Calculate how much we need to rotate the dihedral angle:
        dihedral=dihedral-dihedral_angle([u, v, Origin, D])
        # Calculate the rotation matrix
        Rotation = Rv(Origin-v,np.deg2rad(dihedral))
        # Rotate the dihedral angle by 'dihedral' degrees and translate the vector back:
        return np.squeeze(np.asarray(np.dot(Rotation,(D-v))+v))

    def add_atom(self, name, element, coor):
        index = self._selection[-1]
        if index < len(self.data['record']):
            index = len(self.data['record'])-1
        for attr in self.data:
            if attr == "e":
                setattr(self, '_'+attr, np.append(getattr(self, '_'+attr),
                                                  element))
            elif attr == "atomid":
                setattr(self, '_'+attr, np.append(getattr(self, '_'+attr),
                                                  index+1))
            elif attr == "name":
                setattr(self, '_'+attr, np.append(getattr(self, '_'+attr),
                                                  name))
            elif attr == 'coor':
                setattr(self, '_'+attr, np.append(getattr(self, '_'+attr),
                                                  np.expand_dims(coor, axis=0),
                                                  axis=0))
            else:
                setattr(self, '_'+attr, np.append(getattr(self, '_'+attr),
                                                  getattr(self, attr)[-1]))
        setattr(self, '_selection', np.append(self.__dict__['_selection'],
                                              index+1 ))
        setattr(self, 'natoms', self.natoms+1)

    def reorder(self):
        for idx,atom2 in enumerate(self._rotamers['atoms']+self._rotamers['hydrogens']):
            if self.name[idx] != atom2:
                    idx2 = np.argwhere(self.name == atom2)[0]
                    index = np.ndarray((1,),dtype='int')
                    index[0,]=np.array(self.__dict__['_selection'][0],dtype='int')
                    for attr in ["record", "name", "b", "q", "coor", "resn", "resi","icode", "e", "charge", "chain", "altloc"]:
                        self.__dict__['_'+attr][index+idx],self.__dict__['_'+attr][index+idx2]=self.__dict__['_'+attr][index+idx2],self.__dict__['_'+attr][index+idx]