"""
Module for LIWC related functionality.
"""
import os
from typing import Tuple

import numpy as np
import pandas as pd
import swifter
from tqdm import tqdm
import pysentiment as ps

from connection import MySQLConnector

class LIWCScores:
    """
    Class responsible for the calculation of LIWC analysis scores.
    """
    BASE_PATH = "./out/LIWC/output/"

    # Correlation matrix for signficance level p<0.05.
    corr_05 = np.array([[0, 0, -0.21, 0.11, 0],
                        [0.12, 0, -0.16, 0, 0],
                        [0, 0.11, -0.1, 0.18, 0],
                        [0.1, 0, -0.19, 0.08, 0],
                        [-0.15, 0.16, -0.12, 0, 0],
                        [0, 0, 0, 0, 0],
                        [0.11, 0, -0.13, 0, -0.17],
                        [0, 0, -0.11, 0, -0.09],
                        [-0.11, 0, 0.2, 0, 0.09],
                        [0, 0, 0.17, 0, 0],
                        [0, -0.12, -0.08, 0.11, 0],
                        [0, 0.09, -0.12, 0, 0],
                        [0, 0.1, -0.15, 0.18, 0],
                        [0, 0.11, -0.11, 0.14, 0],
                        [0.08, 0, 0, 0.15, 0.16],
                        [0.16, 0, 0, -0.15, -0.18],
                        [0.17, 0, 0, 0, 0],
                        [0.13, 0, 0, -0.23, -0.19],
                        [0.1, 0, 0, 0, -0.11],
                        [0.13, 0, -0.09, 0, -0.11],
                        [0.11, -0.09, 0, -0.11, -0.12],
                        [0, 0, 0, 0, 0],
                        [0.13, 0, -0.12, 0, -0.13],
                        [0.09, -0.13, 0, 0, 0],
                        [0.12, -0.11, 0, 0, -0.1],
                        [0.13, 0.1, 0, 0, -0.1],
                        [0, 0.09, -0.11, 0, -0.1],
                        [0, 0, 0, 0.09, 0],
                        [0, 0.12, -0.08, 0, -0.12],
                        [0.1, 0, 0, 0.1, 0],
                        [0, 0.15, -0.14, 0.13, 0],
                        [0, 0.13, 0, 0, 0],
                        [-0.08, 0.15, -0.14, 0.15, 0],
                        [-0.08, 0.15, 0, 0.11, 0],
                        [0, 0.09, -0.17, 0.19, 0],
                        [0, 0.13, -0.09, 0, -0.12],
                        [0, 0, -0.22, 0.12, 0.09],
                        [0, 0, -0.16, 0.1, 0],
                        [0, 0, -0.16, 0, 0],
                        [0, 0, 0, 0, 0],
                        [-0.09, 0, -0.11, 0.16, 0],
                        [-0.1, 0.09, -0.15, 0.11, 0.09],
                        [0, 0, -0.11, 0.11, 0],
                        [0, 0.09, 0.11, 0.18, 0],
                        [0.1, 0, 0, 0, -0.16],
                        [0, 0, -0.22, 0.14, 0],
                        [0, -0.12, 0, 0, 0],
                        [0, 0, 0, 0, 0],
                        [0, -0.08, 0, 0, 0],
                        [0, -0.09, 0, 0, 0.14],
                        [0, 0.08, -0.17, 0.15, 0],
                        [0, 0, -0.2, 0.19, 0],
                        [0, 0, -0.14, 0, 0],
                        [0, 0, 0, 0, 0],
                        [0, 0.13, 0, 0.08, -0.11],
                        [0, 0, 0, -0.11, 0],
                        [0, 0, 0, 0, 0],
                        [0, 0.11, 0, 0, 0],
                        [0, 0, 0.15, -0.13, -0.12],
                        [0, 0.14, -0.09, 0.09, 0],
                        [0, 0.1, 0, 0.09, 0],
                        [0, 0.17, 0, 0.08, 0],
                        [0, 0, -0.15, 0, 0],
                        [0.1, 0, -0.14, 0.11, 0],
                        [0, 0, -0.2, 0, 0],
                        [0.11, 0, 0, -0.21, -0.14]])
    # Correlation matrix for signficance level p<0.01.
    corr_01 = np.array([[0, 0, -0.21, 0.11, 0],
                        [0.12, 0, -0.16, 0, 0],
                        [0, 0.11, 0, 0.18, 0],
                        [0, 0, -0.19, 0, 0],
                        [-0.15, 0.16, -0.12, 0, 0],
                        [0, 0, 0, 0, 0],
                        [0.11, 0, -0.13, 0, -0.17],
                        [0, 0, -0.11, 0, 0],
                        [-0.11, 0, 0.2, 0, 0],
                        [0, 0, 0.17, 0, 0],
                        [0, -0.12, 0, 0, 0],
                        [0, 0, -0.12, 0, 0],
                        [0, 0, -0.15, 0.18, 0],
                        [0, 0, -0.11, 0.14, 0],
                        [0, 0, 0, 0.15, 0.16],
                        [0.16, 0, 0, -0.15, -0.18],
                        [0.17, 0, 0, 0, 0],
                        [0.13, 0, 0, -0.23, -0.19],
                        [0, 0, 0, 0, 0],
                        [0.13, 0, 0, 0, -0.11],
                        [0.11, 0, 0, -0.11, -0.12],
                        [0, 0, 0, 0, 0],
                        [0.13, 0, -0.12, 0, -0.13],
                        [0, -0.13, 0, 0, 0],
                        [0.12, 0, 0, 0, 0],
                        [0.13, 0, 0, 0, 0],
                        [0, 0, -0.11, 0, 0],
                        [0, 0, 0, 0, 0],
                        [0, 0.12, 0, 0, -0.12],
                        [0, 0, 0, 0, 0],
                        [0, 0.15, -0.14, 0.13, 0],
                        [0, 0.13, 0, 0, 0],
                        [0, 0.15, -0.14, 0.15, 0],
                        [0, 0.15, 0, 0.11, 0],
                        [0, 0, -0.17, 0.19, 0],
                        [0, 0.13, 0, 0, -0.12],
                        [0, 0, -0.22, 0.12, 0],
                        [0, 0, -0.16, 0, 0],
                        [0, 0, -0.16, 0, 0],
                        [0, 0, 0, 0, 0],
                        [0, 0, -0.11, 0.16, 0],
                        [0, 0, -0.15, 0.11, 0],
                        [0, 0, -0.11, 0.11, 0],
                        [0, 0, 0.11, 0.18, 0],
                        [0, 0, 0, 0, -0.16],
                        [0, 0, -0.22, 0.14, 0],
                        [0, -0.12, 0, 0, 0],
                        [0, 0, 0, 0, 0],
                        [0, 0, 0, 0, 0],
                        [0, 0, 0, 0, 0.14],
                        [0, 0, -0.17, 0.15, 0],
                        [0, 0, -0.2, 0.19, 0],
                        [0, 0, -0.14, 0, 0],
                        [0, 0, 0, 0, 0],
                        [0, 0.13, 0, 0, -0.11],
                        [0, 0, 0, -0.11, 0],
                        [0, 0, 0, 0, 0],
                        [0, 0.11, 0, 0, 0],
                        [0, 0, 0.15, -0.13, -0.12],
                        [0, 0.14, 0, 0, 0],
                        [0, 0, 0, 0, 0],
                        [0, 0.17, 0, 0, 0],
                        [0, 0, -0.15, 0, 0],
                        [0, 0, -0.14, 0.11, 0],
                        [0, 0, -0.2, 0, 0],
                        [0.11, 0, 0, -0.21, -0.14]])
    # Correlation matrix for signficance level p<0.001.
    corr_001 = np.array([[0, 0, -0.21, 0, 0],
                         [0, 0, -0.16, 0, 0],
                         [0, 0, 0, 0.18, 0],
                         [0, 0, -0.19, 0, 0],
                         [-0.15, 0.16, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, -0.17],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0.2, 0, 0],
                         [0, 0, 0.17, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, -0.15, 0.18, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0.15, 0.16],
                         [0.16, 0, 0, -0.15, -0.18],
                         [0.17, 0, 0, 0, 0],
                         [0, 0, 0, -0.23, -0.19],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0.15, -0.14, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0.15, -0.14, 0.15, 0],
                         [0, 0.15, 0, 0, 0],
                         [0, 0, -0.17, 0.19, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, -0.22, 0, 0],
                         [0, 0, -0.16, 0, 0],
                         [0, 0, -0.16, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0.16, 0],
                         [0, 0, -0.15, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0.18, 0],
                         [0, 0, 0, 0, -0.16],
                         [0, 0, -0.22, 0.14, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0.14],
                         [0, 0, -0.17, 0.15, 0],
                         [0, 0, -0.2, 0.19, 0],
                         [0, 0, -0.14, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0, 0.15, 0, 0],
                         [0, 0.14, 0, 0, 0],
                         [0, 0, 0, 0, 0],
                         [0, 0.17, 0, 0, 0],
                         [0, 0, -0.15, 0, 0],
                         [0, 0, -0.14, 0, 0],
                         [0, 0, -0.2, 0, 0],
                         [0, 0, 0, -0.21, -0.14]])

    #    N  E  O  A  C
    corr_05_golbeck = np.array([
        # [0, 0, 0, 0, 0],  # pronoun
        # [0, 0, 0, 0, 0],  # i
        # [0, 0, 0, 0, 0],  # we
        # [0, 0, 0, 0, 0],  # ppron
        [0, 0, 0, 0.364, 0.252],  # you
        # [0, 0, 0, 0, 0],  # shehe
        [0, 0, 0, 0, -0.374],  # negate
        # [0, 0, 0, 0, 0],  # assent
        [0, 0, 0.396, 0, 0],  # article
        # [0, 0, 0, 0, 0],  # preps
        # [0, 0, 0, 0, 0],  # number
        # [0, 0, 0, 0, 0],  # affect
        # [0, 0, 0, 0, 0],  # posemo
        [0, 0, 0, 0, -0.268],  # negemo
        # [0, 0, 0, 0, 0],  # anx
        # [0, 0, 0, 0, 0],  # anger
        [0, 0, 0, 0, -0.253],  # sad
        [0, 0, 0, 0, -0.244],  # cogmech
        [0, 0, 0.264, -0.258, 0],  # cause
        # [0, 0, 0, 0, 0],  # insight
        [0, 0, 0, 0, -0.292],  # discrep
        # [0, 0, 0, 0, 0],  # inhib
        # [0, 0, 0, 0, 0],  # tentat
        [0, 0, 0.347, 0, 0],  # certain
        # [0, 0, 0, 0, 0],  # percept
        # [0, 0, 0, 0, 0],  # see
        [0.335, 0, 0, 0, 0],  # hear
        [0.244, 0, 0, 0, -0.236],  # feel
        [0, 0.262, 0, 0, 0],  # social
        # [0, 0, 0, 0, 0],  # friend
        [0, 0.338, 0, 0, 0],  # family
        [0, 0, 0.251, 0, 0],  # humans
        # [0, 0, 0, 0, 0],  # time
        # [0, 0, 0, 0, 0],  # past
        # [0, 0, 0, 0, 0],  # present
        [0, 0, 0, 0, -0.286],  # future
        # [0, 0, 0, 0, 0],  # space
        # [0, 0, 0, 0, 0],  # incl
        # [0, 0, 0, 0, 0],  # excl
        # [0, 0, 0, 0, 0],  # motion
        [0, 0, 0.426, 0, 0.330],  # work
        [0, 0, 0, -0.240, 0],  # achieve
        # [0, 0, 0, 0, 0],  # leisure
        # [0, 0, 0, 0, 0],  # home
        [0, 0, 0, -0.259, 0],  # money
        [0.383, 0, 0, 0, 0],  # relig
        [0, 0, 0, 0, -0.332],  # death
        [0, 0, -0.239, 0, 0],  # bio
        [0, 0, -0.299, 0, 0],  # body
        # [0, 0, 0, 0, 0],  # sexual
        [0, 0, 0, 0.247, 0],  # ingest
        # [0, 0, 0, 0, 0],  # swear
        [0, 0, 0, 0, -0.284],  # auxverb
        [0, 0, 0.238, 0, 0],  # quant
        [0, -0.277, 0, 0, 0],  # health
        [0, 0, 0, 0, -0.272],  # filler
        [0, 0, 0, 0, -0.24],  # Comma
        [0, 0, 0, 0, 0.322],  # Colon
        [0, 0.263, 0, 0, 0],  # QMark
        [0.317, 0, -0.295, 0, 0.260],  # Exclam
        [0, -0.254, -0.302, 0, 0]])  # Parenth

    def __init__(self, signf=0.05, version=2001, author='yarkoni', feature=None):
        """
        Constructor for creating an LIWC calculation instance.
        :param signf: Significance level (p <0.05, <0.01, or <0.001).
        :param version: LIWC Version (2001, 2007, or 2015).
        :param author: LIWC correlation values author (yarkoni/golbeck).
        :param feature: Feature for which the LIWC is executed
        """
        self.signf = signf
        self.version = version
        self.author = author
        if feature:
            self.path = self.BASE_PATH + 'LIWC{} Results - {}.csv'.format(self.version, feature)
        else:
            self.path = self.BASE_PATH + 'LIWC{} Results.csv'.format(self.version)

    def get_neuroticism(self, cat_scores) -> float:
        """
        Calculate the neuroticism of the user.
        :param cat_scores: LIWC category scores.
        :return: Big Five personality score for Neuroticism.
        """
        return self._get_significance_matrix()[:, 0].dot(cat_scores)

    def get_extraversion(self, cat_scores) -> float:
        """
        Calculate the extraversion of the user.
        :param cat_scores: LIWC category scores.
        :return: Big Five personality score for Extraversion.
        """
        return self._get_significance_matrix()[:, 1].dot(cat_scores)

    def get_openness(self, cat_scores) -> float:
        """
        Calculate the openness of the user.
        :param cat_scores: LIWC category scores.
        :return: Big Five personality score for Openness.
        """
        return self._get_significance_matrix()[:, 2].dot(cat_scores)

    def get_agreeableness(self, cat_scores) -> float:
        """
        Calculate the agreeableness of the user.
        :return: Big Five personality score for Agreeableness.
        """
        return self._get_significance_matrix()[:, 3].dot(cat_scores)

    def get_conscientiousness(self, cat_scores) -> float:
        """
        Calculate the "conscientiousness" of the user.
        :param cat_scores: LIWC category scores.
        :return: Big Five personality score for Conscientiousness.
        """
        return self._get_significance_matrix()[:, 4].dot(cat_scores)

    def get_big_five_scores(self, cat_scores) -> Tuple[float, float, float, float, float]:
        """
        Get all of the Big-Five scores as a single tuple.
        :return: Tuple containing the Big-Five scores in OCEAN order.
        """
        return self.get_openness(cat_scores), self.get_conscientiousness(
            cat_scores), self.get_extraversion(cat_scores), self.get_agreeableness(
            cat_scores), self.get_neuroticism(cat_scores)

    def _get_significance_matrix(self) -> np.array:
        """
        Get the significance level matrix. I.e., get the matrix containing the correlations with
        the right significance level.

        Note that for author=golbeck, the signf variable does not have any influence as only p<0.05
        is available.
        :return: Correlation matrix with the right significance.
        """
        if self.author == 'yarkoni':
            if self.signf == 0.05:
                return self.corr_05
            if self.signf == 0.01:
                return self.corr_01
            return self.corr_001
        return self.corr_05_golbeck  # Author golbeck

    @staticmethod
    def _get_columns(version, author):
        """
        Get the columns that below to the given LIWC version. Different LIWC versions have different
        column names and different columns available.
        :param version: LIWC version for which columns are selected.
        :param author: Author versino for which the columns are selected.
        :return: Columns available for the given LIWC version.
        """
        # Version 2001
        if version == 2001:
            return ['Pronoun', 'I', 'We', 'Self', 'You', 'Other', 'Negate', 'Assent', 'Article',
                    'Preps', 'Number', 'Affect', 'Posemo', 'Posfeel', 'Optim', 'Negemo', 'Anx',
                    'Anger', 'Sad', 'Cogmech', 'Cause', 'Insight', 'Discrep', 'Inhib', 'Tentat',
                    'Certain', 'Senses', 'See', 'Hear', 'Feel', 'Social', 'Comm', 'Othref',
                    'Friends', 'Family', 'Humans', 'Time', 'Past', 'Present', 'Future', 'Space',
                    'Up', 'Down', 'Incl', 'Excl', 'Motion', 'Occup', 'School', 'Job', 'Achieve',
                    'Leisure', 'Home', 'Sports', 'TV', 'Music', 'Money', 'Metaph', 'Relig',
                    'Death',
                    'Physcal', 'Body', 'Sexual', 'Eating', 'Sleep', 'Groom', 'Swear']

        # Version 2007
        if version == 2007 and author == 'golbeck':
            return ['you', 'negate', 'article', 'negemo', 'sad', 'cogmech', 'cause', 'discrep',
                    'certain', 'hear', 'feel', 'social', 'family', 'humans', 'future', 'work',
                    'achieve', 'money', 'relig', 'death', 'bio', 'body', 'ingest', 'auxverb',
                    'quant', 'health', 'filler', 'Comma', 'Colon', 'QMark', 'Exclam', 'Parenth']

        # Version 2015 / Different author combination
        return [
            'pronoun', 'i', 'we', 'ppron', 'you', 'shehe', 'negate', 'assent', 'article',
            'preps', 'number', 'affect', 'posemo', 'negemo', 'anx', 'anger', 'sad', 'cogmech',
            'cause', 'insight', 'discrep', 'inhib', 'tentat', 'certain', 'percept', 'see',
            'hear', 'feel', 'social', 'friend', 'family', 'humans', 'time', 'past', 'present',
            'future', 'space', 'incl', 'excl', 'motion', 'work', 'achieve', 'leisure', 'home',
            'money', 'relig', 'death', 'bio', 'body', 'sexual', 'ingest', 'swear', 'auxverb',
            'quant', 'health', 'filler', 'Comma', 'Colon', 'QMark', 'Exclam', 'Parenth']

    def _get_liwc_df(self) -> pd.DataFrame:
        """
        Get the liwc scores as a dataframe.
        :return: LIWC scores in a pandas dataframe.
        """
        data = pd.read_csv(self.path)
        data.index = pd.to_numeric(data['Filename'].str.rstrip('.txt'))
        return data

    def _get_category_scores(self, data) -> pd.DataFrame:
        """
        Get the category scores from the LIWC CSV file and convert to a pandas dataframe.
        :return: LIWC category scores in a pandas dataframe.
        """
        return data[self._get_columns(self.version, self.author)]

    def apply_corrections(self, liwc_data, liwc_scores):
        """
        Apply corrections on the found scores applicable to the GitHub data.
        :param liwc_data: LIWC data containing the category scores.
        :param liwc_scores: Scores already found from the LIWC data.
        :return Updated dataframe with LIWC scores corrected for software engineering.
        """
        print('Applying corrections..')
        if self.author == 'yarkoni' and self.version == 2001:
            liwc_data_copy = liwc_data.copy()
            liwc_data_copy['Pronoun'] = np.log(liwc_data_copy['Pronoun'] + 1)
            liwc_data_copy['Social'] = liwc_data_copy['Social'] ** (1 / 2)
            liwc_scores['Agreeableness'] = liwc_data_copy.swifter.apply(
                self.get_agreeableness, axis=1, result_type='expand'
            )
            liwc_data_copy = liwc_data.copy()
            liwc_data_copy['Social'] = liwc_data_copy['Social'] ** (1 / 2)
            liwc_scores['Extraversion'] = liwc_data_copy.swifter.apply(
                self.get_extraversion, axis=1, result_type='expand'
            )
        elif self.author == 'golbeck' and self.version == 2007:
            liwc_data_copy = liwc_data.copy()
            liwc_data_copy['Exclam'] = np.log(liwc_data_copy['Exclam'] + 1)
            liwc_data_copy['hear'] = np.log(liwc_data_copy['hear'] + 1)
            liwc_scores['Neuroticism'] = liwc_data_copy.swifter.apply(
                self.get_neuroticism, axis=1, result_type='expand'
            )
            liwc_data_copy = liwc_data.copy()
            liwc_data_copy['you'] = liwc_data_copy['you'] ** (1 / 3)
            liwc_scores['Agreeableness'] = liwc_data_copy.swifter.apply(
                self.get_agreeableness, axis=1, result_type='expand'
            )
        return liwc_scores

    @staticmethod
    def calculate_sentiment(positive, negative):
        """
        Calculate the sentiment scores.
        :param positive: Number of words with positive sentiment.
        :param negative: Number of words with negative sentiment.
        :return: Average sentiment score multiplied by the correlation coefficient of Golbeck.
        """
        denominator = (positive - negative)
        numerator = (positive + negative)
        if numerator == 0:
            return 0
        return 0.268 * (denominator / numerator)

    def add_sentiment_scores(self, liwc_scores):
        """
        Add the sentment scores to the Openness score of Golbeck.
        :param liwc_scores: The LIWC scores found for Golbeck.
        :return: LIWC Scores with the Openness updated with the sentiment scores.
        """
        print("Applying sentiment scores.")
        hiv4 = ps.HIV4()
        for idx in tqdm(liwc_scores.index):
            try:
                with open(os.environ['LIWC_INPUT'] + str(idx) + '.txt', 'r') as f:
                    finput = f.readlines()[0]
                    tokens = hiv4.tokenize(finput)
                    sentiment_scores = hiv4.get_score(tokens)
                    liwc_scores.at[idx, 'Openness'] = liwc_scores.at[idx, 'Openness'] + \
                                                    self.calculate_sentiment(
                                                        sentiment_scores['Positive'],
                                                        sentiment_scores['Negative'])
            except FileNotFoundError:
                print('User {} does not exist in the {} {} results csv'.format(idx, self.author,
                                                                               self.version))
        return liwc_scores

    def get_liwc_scores(self) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        Get the LIWC personality scores calculated from the CSVs generated for each person.
        :return: Tuple of raw scores and normalized scores.
        """
        big_five_cols = ['Openness', 'Conscientiousness', 'Extraversion',
                         'Agreeableness', 'Neuroticism']

        liwc_data = self._get_category_scores(self._get_liwc_df())
        liwc_scores = pd.DataFrame(columns=big_five_cols)
        liwc_scores[big_five_cols] = liwc_data.swifter.apply(
            self.get_big_five_scores, axis=1, result_type='expand'
        )
        liwc_scores = self.apply_corrections(liwc_data, liwc_scores)
        if self.author == 'golbeck':
            liwc_scores = self.add_sentiment_scores(liwc_scores)

        # Min-max normalization
        liwc_scores_normalized = (liwc_scores - liwc_scores.min()) / (
                liwc_scores.max() - liwc_scores.min()) * 100
        return liwc_scores.round(15), liwc_scores_normalized.round(15)


#!/usr/bin/env python3

"""
TODO ligand.py Refactor:
    DO:
        Move module-specific methods such as openmm_coordinates(); read_tdrive(); read_geometric_traj
            to their relevant files/classes
        Fix naming; consistency wrt get/find; clarity on all of the dihedral variables
            (what is dih_start, how is it different to di_starts etc)
        Perform checks after reading input (check_names_are_unique(), validate_info(), etc)
    CONSIDER:
        Remove / replace DefaultsMixin with inheritance, dict or some other solution
        Remove any repeated or unnecessary variables
            Should state be handled in ligand or run?
        Change the structure and type of some variables for clarity
            Do we access via index too often; should we use e.g. SimpleNamespaces/NamedTupleS?
        Be more strict about public/private class/method/function naming?
"""

import decimal
import os
import xml.etree.ElementTree as ET
from datetime import datetime
from typing import Any, Dict, List, Optional, Set, Tuple, Union
from xml.dom.minidom import parseString

import networkx as nx
import numpy as np
import qcelemental as qcel
from pydantic import Field, validator
from qcelemental.models.types import Array
from rdkit import Chem

try:
    # fix for the openmm namechange
    from openmm import unit
    from openmm.app import Aromatic, Double, PDBFile, Single, Topology, Triple
    from openmm.app.element import Element
except (ModuleNotFoundError, ImportError):
    from simtk import unit
    from simtk.openmm.app import Aromatic, Double, Single, Topology, Triple, PDBFile
    from simtk.openmm.app.element import Element

import qubekit
from qubekit.forcefield import (
    BaseForceGroup,
    HarmonicAngleForce,
    HarmonicBondForce,
    LennardJones126Force,
    PeriodicImproperTorsionForce,
    PeriodicTorsionForce,
    RBImproperTorsionForce,
    RBProperTorsionForce,
    VirtualSiteGroup,
)
from qubekit.molecules.components import Atom, Bond, TorsionDriveData
from qubekit.molecules.utils import RDKit, ReadInput
from qubekit.utils import constants
from qubekit.utils.datastructures import SchemaBase
from qubekit.utils.exceptions import (
    ConformerError,
    FileTypeError,
    MissingReferenceData,
    StereoChemistryError,
    TopologyMismatch,
)

class Molecule(SchemaBase):
    """Base class for ligands and proteins.

    The class is a simple representation of the molecule as a list of atom and bond objects, many attributes are then
    inferred from these core objects.
    """

    atoms: List[Atom] = Field(
        ..., description="A list of QUBEKit atom objects which make up the molecule."
    )
    bonds: Optional[List[Bond]] = Field(
        None,
        description="The list of QUBEKit bond objects which connect the individual atoms.",
    )
    coordinates: Optional[Array[float]] = Field(
        None,
        description="A numpy arrary of the current cartesian positions of each atom in angstrom, this must be of size (n_atoms, 3)",
    )
    multiplicity: int = Field(
        1,
        description="The integer multiplicity of the molecule which is used in QM calculations.",
    )
    name: str = Field(
        "unk",
        description="An optional name string which will be used in all file IO calls by default.",
    )
    provenance: Dict[str, Any] = Field(
        dict(creator="QUBEKit", version=qubekit.__version__),
        description="Information on the version and method used to create this molecule.",
    )
    extra_sites: VirtualSiteGroup = Field(
        VirtualSiteGroup(),
        description="A force object which records any virtual sistes in the molecule",
    )
    BondForce: Union[HarmonicBondForce] = Field(
        HarmonicBondForce(),
        description="A force object which records bonded interactions between pairs of atoms",
    )
    AngleForce: Union[HarmonicAngleForce] = Field(
        HarmonicAngleForce(),
        description="A force object which records angle interactions between atom triplets.",
    )
    TorsionForce: PeriodicTorsionForce = Field(
        PeriodicTorsionForce(),
        description="A force object which records torsion interactions between atom quartets, using a periodic function.",
    )
    ImproperTorsionForce: PeriodicImproperTorsionForce = Field(
        PeriodicImproperTorsionForce(),
        description="A force group which records improper torsion interactions between atom quartets using a periodic function.",
    )
    RBTorsionForce: RBProperTorsionForce = Field(
        RBProperTorsionForce(),
        description="A force group which records torsion interactions between atom quartets using a RB function.",
    )
    ImproperRBTorsionForce: RBImproperTorsionForce = Field(
        RBImproperTorsionForce(),
        description="A force object which records improper torsion interactions between atom quartets using a RB function.",
    )
    NonbondedForce: Union[LennardJones126Force] = Field(
        LennardJones126Force(),
        description="A force group which records atom nonbonded parameters.",
    )
    chargemol_coords: Optional[Array[float]] = Field(
        None,
        description="The coordinates used to calculate the chargemol quantities, "
        "this is a reorientated conformation",
    )

    @validator("coordinates", "chargemol_coords")
    def _reshape_coords(cls, coordinates: Optional[np.ndarray]) -> Optional[np.ndarray]:
        if coordinates is not None:
            return coordinates.reshape((-1, 3))
        else:
            return coordinates

    def __init__(
        self,
        atoms: List[Atom],
        bonds: Optional[List[Bond]] = None,
        coordinates: Optional[np.ndarray] = None,
        multiplicity: int = 1,
        name: str = "unk",
        routine: Optional[Set] = None,
        provenance: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        """
        Init the molecule using the basic information.

        Note:
            This method updates the provanance info.

        Args:
            atoms:
                A list of QUBEKit atom objects in the molecule.
            bonds:
                A list of QUBEKit bond objects in the molecule.
            coordinates:
                A numpy array of the current cartesian positions of each atom, this must be of size (n_atoms, 3)
            multiplicity:
                The integer multiplicity of the molecule which is used in QM calculations.
            name:
                An optional name string which will be used in all file IO calls by default.
            routine:
                The set of strings which encode the routine information used to create the molecule.

            bool; is the current execution starting from the beginning (False) or restarting (True)?
        """
        # the way the molecule was made
        method = routine or {"__init__"}
        if provenance is None:
            new_provenance = dict(
                creator="QUBEKit", version=qubekit.__version__, routine=method
            )
        else:
            # make sure we respect the provenance when parsing a json file
            new_provenance = provenance
            new_provenance["routine"] = set(new_provenance["routine"])

        super(Molecule, self).__init__(
            atoms=atoms,
            bonds=bonds,
            multiplicity=multiplicity,
            name=name,
            coordinates=coordinates,
            provenance=new_provenance,
            **kwargs,
        )

    def to_topology(self) -> nx.Graph:
        """
        Build a networkx representation of the molecule.
        TODO add other attributes to the graph?
        """
        graph = nx.Graph()
        for atom in self.atoms:
            graph.add_node(atom.atom_index)

        for bond in self.bonds:
            graph.add_edge(bond.atom1_index, bond.atom2_index)
        return graph

    def to_file(self, file_name: str) -> None:
        """
        Write the molecule object to file working out the file type from the extension.
        Works with PDB, MOL, SDF, XYZ, Json any other we want?
        """
        if ".json" in file_name:
            with open(file_name, "w") as output:
                output.write(self.json(indent=2))
        else:
            return RDKit.mol_to_file(rdkit_mol=self.to_rdkit(), file_name=file_name)

    def generate_conformers(self, n_conformers: int) -> List[np.ndarray]:
        """
        Generate a list of conformers and return them this includes the input conformer

        Args:
            n_conformers: The number of conformers which should be generated.
        """
        rd_mol = self.to_rdkit()
        return RDKit.generate_conformers(rdkit_mol=rd_mol, conformer_no=n_conformers)

    def to_multiconformer_file(
        self, file_name: str, positions: List[np.ndarray]
    ) -> None:
        """
        Write the molecule to a file allowing multipule conformers.

        As the ligand object only holds one set of coordinates at once a list of coords can be passed here to allow
        multiconformer support.

        Args:
            file_name:
                The name of the file that should be created, the type is inferred by the suffix.
            positions:
                A list of Cartesian coordinates of shape (n_atoms, 3).
        """
        rd_mol = self.to_rdkit()
        rd_mol.RemoveAllConformers()
        # add the conformers
        if not isinstance(positions, list):
            positions = [
                positions,
            ]
        for conformer in positions:
            RDKit.add_conformer(rdkit_mol=rd_mol, conformer_coordinates=conformer)
        return RDKit.mol_to_multiconformer_file(rdkit_mol=rd_mol, file_name=file_name)

    def get_atom_with_name(self, name):
        """
        Search through the molecule for an atom with that name and return it when found
        :param name: The name of the atom we are looking for
        :return: The QUBE Atom object with the name
        """

        for atom in self.atoms:
            if atom.atom_name == name:
                return atom
        raise AttributeError("No atom found with that name.")

    def get_bond_between(self, atom1_index: int, atom2_index: int) -> Bond:
        """
        Try and find a bond between the two atom indices.

        The bond may not have the atoms in the expected order.

        Args:
            atom1_index:
                The index of the first atom in the atoms list.
            atom2_index:
                The index of the second atom in the atoms list.

        Returns:
            The bond object between the two target atoms.

        Raises:
            TopologyMismatch:
                When no bond can be found between the atoms.
        """
        target = [atom1_index, atom2_index]
        for bond in self.bonds:
            if bond.atom1_index in target and bond.atom2_index in target:
                return bond
        raise TopologyMismatch(
            f"There is no bond between atoms {atom1_index} and {atom2_index} in this molecule."
        )

    @property
    def has_unique_atom_names(self) -> bool:
        """
        Check if the molecule has unique atom names or not this will help with pdb file writing.
        """
        atom_names = set([atom.atom_name for atom in self.atoms])
        if len(atom_names) == self.n_atoms:
            return True
        return False

    @property
    def improper_torsions(self) -> Optional[List[Tuple[int, int, int, int]]]:
        """A list of improper atom tuples where the first atom is central."""

        improper_torsions = []
        topology = self.to_topology()
        for node in topology.nodes:
            near = sorted(list(nx.neighbors(topology, node)))
            # if the atom has 3 bonds it could be an improper
            # Check if an sp2 carbon or N
            if len(near) == 3 and (
                self.atoms[node].atomic_symbol == "C"
                or self.atoms[node].atomic_symbol == "N"
            ):
                # Store each combination of the improper torsion
                improper_torsions.append((node, near[0], near[1], near[2]))
        return improper_torsions or None

    @property
    def n_improper_torsions(self) -> int:
        """The number of unique improper torsions."""
        impropers = self.improper_torsions
        if impropers is None:
            return 0
        return len(impropers)

    @property
    def angles(self) -> Optional[List[Tuple[int, int, int]]]:
        """A List of angles from the topology."""

        angles = []
        topology = self.to_topology()
        for node in topology.nodes:
            bonded = sorted(list(nx.neighbors(topology, node)))

            # Check that the atom has more than one bond
            if len(bonded) < 2:
                continue

            # Find all possible angle combinations from the list
            for i in range(len(bonded)):
                for j in range(i + 1, len(bonded)):
                    atom1, atom3 = bonded[i], bonded[j]
                    angles.append((atom1, node, atom3))
        return angles or None

    @property
    def charge(self) -> int:
        """
        Return the integer charge of the molecule as the sum of the formal charge.
        """
        return sum([atom.formal_charge for atom in self.atoms])

    @property
    def n_angles(self) -> int:
        """The number of angles in the molecule."""
        angles = self.angles
        if angles is None:
            return 0
        return len(angles)

    def measure_bonds(self) -> Dict[Tuple[int, int], float]:
        """
        Find the length of all bonds in the molecule for the given conformer in  angstroms.

        Returns:
            A dictionary of the bond lengths stored by bond tuple.
        """

        bond_lengths = {}

        for bond in self.bonds:
            atom1 = self.coordinates[bond.atom1_index]
            atom2 = self.coordinates[bond.atom2_index]
            edge = (bond.atom1_index, bond.atom2_index)
            bond_lengths[edge] = np.linalg.norm(atom2 - atom1)

        return bond_lengths

    @property
    def n_bonds(self) -> int:
        """The number of bonds in the topology."""
        bonds = self.bonds
        if bonds is None:
            return 0
        return len(bonds)

    @property
    def dihedrals(
        self,
    ) -> Optional[Dict[Tuple[int, int], List[Tuple[int, int, int, int]]]]:
        """A list of all possible dihedrals that can be found in the topology."""

        dihedrals = {}
        topology = self.to_topology()
        # Work through the network using each edge as a central dihedral bond
        for edge in topology.edges:

            for start in list(nx.neighbors(topology, edge[0])):

                # Check atom not in main bond
                if start != edge[0] and start != edge[1]:

                    for end in list(nx.neighbors(topology, edge[1])):

                        # Check atom not in main bond
                        if end != edge[0] and end != edge[1]:

                            if edge not in dihedrals:
                                # Add the central edge as a key the first time it is used
                                dihedrals[edge] = [(start, edge[0], edge[1], end)]

                            else:
                                # Add the tuple to the correct key.
                                dihedrals[edge].append((start, edge[0], edge[1], end))

        return dihedrals or None

    @property
    def n_dihedrals(self) -> int:
        """The total number of dihedrals in the molecule."""
        dihedrals = self.dihedrals
        if dihedrals is None:
            return 0
        return sum([len(torsions) for torsions in dihedrals.values()])

    def find_rotatable_bonds(
        self, smirks_to_remove: Optional[List[str]] = None
    ) -> Optional[List[Bond]]:
        """
        Args:
            smirks_to_remove:
                Optional list of smirks patterns which will be discarded
                from the rotatable bonds
        Find all rotatable bonds in the molecule.
        Remove any groups which are not relevant for torsion scans.
            e.g. methyl / amine groups
        return:
            The rotatable bonds in the molecule to be used for torsion scans.
        """

        rotatable_bond_smarts = "[!$(*#*)&!D1:1]-&!@[!$(*#*)&!D1:2]"

        rotatable_matches = self.get_smarts_matches(rotatable_bond_smarts)
        if rotatable_matches is None:
            return None

        if smirks_to_remove is not None:
            for smirk in smirks_to_remove:
                matches_to_remove = self.get_smarts_matches(smirk)
                if matches_to_remove is not None:
                    for match in matches_to_remove:
                        try:
                            rotatable_matches.remove(match)
                        except ValueError:
                            try:
                                # If the match is not in the list, it may be in backwards
                                rotatable_matches.remove(tuple(reversed(match)))
                            except ValueError:
                                continue

        # gather a list of bond instances to return
        rotatable_bonds = [self.get_bond_between(*bond) for bond in rotatable_matches]

        return rotatable_bonds or None

    @property
    def n_rotatable_bonds(self) -> int:
        """The number of rotatable bonds."""
        rotatable_bonds = self.find_rotatable_bonds()
        if rotatable_bonds is None:
            return 0
        return len(rotatable_bonds)

    def symmetrise_nonbonded_parameters(self) -> bool:
        """
        Symmetrise all non-bonded force group parameters.

        Using the CIP rankings from RDKit apply symmetry to the non-bonded force group.

        Important:
            We respect the predefined parameters in the non-bonded force group which can be symmetrised.
        """
        # group atom types as they are in a different format to other types
        atom_types = {}
        for atom_index, cip_type in self.atom_types.items():
            atom_types.setdefault(cip_type, []).append((atom_index,))
        for atoms in atom_types.items():
            self._symmetrise_parameters(
                force_group=self.NonbondedForce, parameter_keys=atoms
            )

        return True

    def symmetrise_bonded_parameters(self) -> bool:
        """
        Symmetrise all bond and angle force group parameters.

        Using the CIP rankings from RDKit apply symmetry to the bond and angle force groups.

        Important:
            We respect the predefined parameters in the bond/angle force group which can be symmetrised.
        """

        for bonds in self.bond_types.values():
            self._symmetrise_parameters(
                force_group=self.BondForce, parameter_keys=bonds
            )

        for angles in self.angle_types.values():
            self._symmetrise_parameters(
                force_group=self.AngleForce, parameter_keys=angles
            )

        return True

    def _symmetrise_parameters(
        self, force_group: BaseForceGroup, parameter_keys: List[Tuple[int, ...]]
    ):
        """
        Internal method which applies symmetry to a group of parameter references in a particular force group.

        Args:
            force_group: The force group we should query for parameters.
            parameter_keys: The list of atom indices tuples that the symmetry should be applied to.
        """

        symmetry_attrs = force_group.symmetry_parameters()

        raw_parameter_values = {}
        for parameter_key in parameter_keys:
            param = force_group[parameter_key]
            for attr in symmetry_attrs:
                raw_parameter_values.setdefault(attr, []).append(getattr(param, attr))

        # now average the raw values
        for key, value in raw_parameter_values.items():
            raw_parameter_values[key] = np.array(value).mean()

        # now set back
        for parameter_key in parameter_keys:
            force_group.create_parameter(atoms=parameter_key, **raw_parameter_values)

    def measure_dihedrals(self) -> Optional[Dict[Tuple[int, int, int, int], float]]:
        """
        For the given conformation measure the dihedrals in the topology in degrees.
        """
        dihedrals = self.dihedrals
        if dihedrals is None:
            return None

        dih_phis = {}

        for val in dihedrals.values():
            for torsion in val:
                # Calculate the dihedral angle in the molecule using the molecule data array.
                x1, x2, x3, x4 = [self.coordinates[torsion[i]] for i in range(4)]
                b1, b2, b3 = x2 - x1, x3 - x2, x4 - x3
                t1 = np.linalg.norm(b2) * np.dot(b1, np.cross(b2, b3))
                t2 = np.dot(np.cross(b1, b2), np.cross(b2, b3))
                dih_phis[torsion] = np.degrees(np.arctan2(t1, t2))

        return dih_phis

    def measure_angles(self) -> Optional[Dict[Tuple[int, int, int], float]]:
        """
        For the given conformation measure the angles in the topology in degrees.
        """
        angles = self.angles
        if angles is None:
            return None

        angle_values = {}

        for angle in angles:
            x1 = self.coordinates[angle[0]]
            x2 = self.coordinates[angle[1]]
            x3 = self.coordinates[angle[2]]
            b1, b2 = x1 - x2, x3 - x2
            cosine_angle = np.dot(b1, b2) / (np.linalg.norm(b1) * np.linalg.norm(b2))
            angle_values[angle] = np.degrees(np.arccos(cosine_angle))

        return angle_values

    @property
    def n_atoms(self) -> int:
        """
        Calculate the number of atoms.
        """
        return len(self.atoms)

    def write_parameters(self, file_name: str):
        """
        Take the molecule's parameter set and write an xml file for the molecule.
        """

        tree = self._build_forcefield().getroot()
        messy = ET.tostring(tree, "utf-8")

        pretty_xml_as_string = parseString(messy).toprettyxml(indent="")

        with open(file_name, "w") as xml_doc:
            xml_doc.write(pretty_xml_as_string)

    def _build_forcefield(self):
        """
        Separates the parameters and builds an xml tree ready to be used.

        TODO how do we support OPLS combination rules.
        Important:
            The ordering here should not be changed due to the way sites have to be added.
        """

        # Create XML layout
        root = ET.Element("ForceField")

        ET.SubElement(
            root,
            "QUBEKit",
            attrib={
                "Version": qubekit.__version__,
                "Date": datetime.now().strftime("%Y_%m_%d"),
            },
        )
        AtomTypes = ET.SubElement(root, "AtomTypes")
        Residues = ET.SubElement(root, "Residues")

        resname = "QUP" if self.__class__.__name__ == "Protein" else "MOL"
        Residue = ET.SubElement(Residues, "Residue", name=resname)
        # declare atom `types` and properties
        for atom in self.atoms:
            atom_type = f"QUBE_{atom.atom_index}"
            ET.SubElement(
                AtomTypes,
                "Type",
                attrib={
                    "name": atom_type,
                    "class": str(atom.atom_index),
                    "element": atom.atomic_symbol,
                    "mass": str(atom.atomic_mass),
                },
            )

            ET.SubElement(
                Residue, "Atom", attrib={"name": atom.atom_name, "type": atom_type}
            )

        # add sites to Atomtypes, topology and nonbonded
        for i, site in enumerate(self.extra_sites, start=1):
            site_name = f"v-site{i}"
            site_class = f"X{i}"
            ET.SubElement(
                AtomTypes,
                "Type",
                attrib={"name": site_name, "class": site_class, "mass": "0"},
            )
            # for some reason we swap name and class here but it works !
            ET.SubElement(
                Residue, "Atom", attrib={"name": site_class, "type": site_name}
            )

        BondForce = ET.SubElement(
            root, self.BondForce.openmm_group(), attrib=self.BondForce.xml_data()
        )
        for parameter in self.BondForce:
            ET.SubElement(
                BondForce, parameter.openmm_type(), attrib=parameter.xml_data()
            )
            ET.SubElement(
                Residue,
                "Bond",
                attrib={"from": str(parameter.atoms[0]), "to": str(parameter.atoms[1])},
            )
        AngleForce = ET.SubElement(
            root, self.AngleForce.openmm_group(), attrib=self.AngleForce.xml_data()
        )
        for parameter in self.AngleForce:
            ET.SubElement(
                AngleForce, parameter.openmm_type(), attrib=parameter.xml_data()
            )
        if (
            self.TorsionForce.n_parameters > 0
            or self.ImproperTorsionForce.n_parameters > 0
        ):
            TorsionForce = ET.SubElement(
                root,
                self.TorsionForce.openmm_group(),
                attrib=self.TorsionForce.xml_data(),
            )
            for parameter in self.TorsionForce:
                ET.SubElement(
                    TorsionForce, parameter.openmm_type(), attrib=parameter.xml_data()
                )
            for parameter in self.ImproperTorsionForce:
                ET.SubElement(
                    TorsionForce, parameter.openmm_type(), attrib=parameter.xml_data()
                )
        if (
            self.RBTorsionForce.n_parameters > 0
            or self.ImproperRBTorsionForce.n_parameters > 0
        ):
            RBTorsion = ET.SubElement(
                root,
                self.RBTorsionForce.openmm_group(),
                attrib=self.RBTorsionForce.xml_data(),
            )
            for parameter in self.RBTorsionForce:
                ET.SubElement(
                    RBTorsion, parameter.openmm_type(), attrib=parameter.xml_data()
                )
            for parameter in self.ImproperRBTorsionForce:
                ET.SubElement(
                    RBTorsion, parameter.openmm_type(), attrib=parameter.xml_data()
                )

        # now we add more site info after general bonding
        for i, site in enumerate(self.extra_sites):
            site_data = site.xml_data()
            # we have to add its global index
            site_data["index"] = str(i + self.n_atoms)
            ET.SubElement(Residue, site.openmm_type(), attrib=site_data)

        NonbondedForce = ET.SubElement(
            root,
            self.NonbondedForce.openmm_group(),
            attrib=self.NonbondedForce.xml_data(),
        )
        for parameter in self.NonbondedForce:
            ET.SubElement(
                NonbondedForce, parameter.openmm_type(), attrib=parameter.xml_data()
            )

        for i, site in enumerate(self.extra_sites, start=1):
            site_name = f"v-site{i}"
            ET.SubElement(
                NonbondedForce,
                "Atom",
                attrib={
                    "charge": str(site.charge),
                    "epsilon": "0",
                    "sigma": "1",
                    "type": site_name,
                },
            )

        return ET.ElementTree(root)

    @property
    def bond_types(self) -> Dict[str, List[Tuple[int, int]]]:
        """
        Using the symmetry dict, give each bond a code. If any codes match, the bonds can be symmetrised.
        e.g. bond_symmetry_classes = {(0, 3): '2-0', (0, 4): '2-0', (0, 5): '2-0' ...}
        all of the above bonds (tuples) are of the same type (methyl H-C bonds in same region)
        This dict is then used to produce bond_types.
        bond_types is just a dict where the keys are the string code from above and the values are all
        of the bonds with that particular type.
        """
        atom_types = self.atom_types
        bond_symmetry_classes = {}
        for bond in self.bonds:
            bond_symmetry_classes[(bond.atom1_index, bond.atom2_index)] = (
                f"{atom_types[bond.atom1_index]}-" f"{atom_types[bond.atom2_index]}"
            )

        bond_types = {}
        for key, val in bond_symmetry_classes.items():
            bond_types.setdefault(val, []).append(key)

        bond_types = self._cluster_types(bond_types)
        return bond_types

    @property
    def angle_types(self) -> Dict[str, List[Tuple[int, int, int]]]:
        """
        Using the symmetry dict, give each angle a code. If any codes match, the angles can be symmetrised.
        e.g. angle_symmetry_classes = {(1, 0, 3): '3-2-0', (1, 0, 4): '3-2-0', (1, 0, 5): '3-2-0' ...}
        all of the above angles (tuples) are of the same type (methyl H-C-H angles in same region)
        angle_types is just a dict where the keys are the string code from the above and the values are all
        of the angles with that particular type.
        """
        atom_types = self.atom_types
        angle_symmetry_classes = {}
        for angle in self.angles:
            angle_symmetry_classes[angle] = (
                f"{atom_types[angle[0]]}-"
                f"{atom_types[angle[1]]}-"
                f"{atom_types[angle[2]]}"
            )

        angle_types = {}
        for key, val in angle_symmetry_classes.items():
            angle_types.setdefault(val, []).append(key)

        angle_types = self._cluster_types(angle_types)
        return angle_types

    @property
    def dihedral_types(self) -> Dict[str, List[Tuple[int, int, int, int]]]:
        """
        Using the symmetry dict, give each dihedral a code. If any codes match, the dihedrals can be clustered and their
        parameters should be the same, this is to be used in dihedral fitting so all symmetry equivalent dihedrals are
        optimised at the same time. dihedral_equiv_classes = {(0, 1, 2 ,3): '1-1-2-1'...} all of the tuples are the
        dihedrals index by topology and the strings are the symmetry equivalent atom combinations.
        """
        atom_types = self.atom_types
        dihedral_symmetry_classes = {}
        for dihedral_set in self.dihedrals.values():
            for dihedral in dihedral_set:
                dihedral_symmetry_classes[tuple(dihedral)] = (
                    f"{atom_types[dihedral[0]]}-"
                    f"{atom_types[dihedral[1]]}-"
                    f"{atom_types[dihedral[2]]}-"
                    f"{atom_types[dihedral[3]]}"
                )

        dihedral_types = {}
        for key, val in dihedral_symmetry_classes.items():
            dihedral_types.setdefault(val, []).append(key)

        dihedral_types = self._cluster_types(dihedral_types)
        return dihedral_types

    @property
    def improper_types(self) -> Dict[str, List[Tuple[int, int, int, int]]]:
        """Using the atom symmetry types work out the improper types."""

        atom_types = self.atom_types
        improper_symmetry_classes = {}
        for dihedral in self.improper_torsions:
            improper_symmetry_classes[tuple(dihedral)] = (
                f"{atom_types[dihedral[0]]}-"
                f"{atom_types[dihedral[1]]}-"
                f"{atom_types[dihedral[2]]}-"
                f"{atom_types[dihedral[3]]}"
            )

        improper_types = {}
        for key, val in improper_symmetry_classes.items():
            improper_types.setdefault(val, []).append(key)

        improper_types = self._cluster_types(improper_types)
        return improper_types

    @staticmethod
    def _cluster_types(equiv_classes):
        """
        Function that helps the bond angle and dihedral class finders in clustering the types based on the forward and
        backward type strings.
        :return: clustered equiv class
        """

        new_classes = {}
        for key, item in equiv_classes.items():
            try:
                new_classes[key].extend(item)
            except KeyError:
                try:
                    new_classes[key[::-1]].extend(item)
                except KeyError:
                    new_classes[key] = item

        return new_classes

    @property
    def atom_types(self) -> Dict[int, str]:
        """Returns a dictionary of atom indices mapped to their class or None if there is no rdkit molecule."""

        return RDKit.find_symmetry_classes(self.to_rdkit())

    def to_rdkit(self) -> Chem.Mol:
        """
        Generate an rdkit representation of the QUBEKit ligand object.

        Here we build the molecule and assign the stereochemistry using the coordinates as we should always have a set of coordinates in the model.
        This allows us to skip complicated local vs global stereo chemistry checks however this could break in future.

        Returns:
            An rdkit representation of the molecule.
        """
        from qubekit.utils.helpers import _assert_wrapper

        # TODO what properties should be put in the rdkit molecule? Multiplicity?
        # make an editable molecule
        rd_mol = Chem.RWMol()
        if self.name is not None:
            rd_mol.SetProp("_Name", self.name)

        # when building the molecule we have to loop multiple times
        # so always make sure the indexing is the same in qube and rdkit
        for atom in self.atoms:
            rd_index = rd_mol.AddAtom(atom.to_rdkit())
            assert rd_index == atom.atom_index

        # now we need to add each bond, can not make a bond from python currently
        for bond in self.bonds:
            rd_mol.AddBond(*bond.indices)
            # now get the bond back to edit it
            rd_bond: Chem.Bond = rd_mol.GetBondBetweenAtoms(*bond.indices)
            rd_bond.SetIsAromatic(bond.aromatic)
            rd_bond.SetBondType(bond.rdkit_type)

        Chem.SanitizeMol(
            rd_mol,
            Chem.SANITIZE_ALL ^ Chem.SANITIZE_ADJUSTHS ^ Chem.SANITIZE_SETAROMATICITY,
        )
        # must use openff MDL model for compatibility
        Chem.SetAromaticity(rd_mol, Chem.AromaticityModel.AROMATICITY_MDL)

        # conformers
        rd_mol = RDKit.add_conformer(
            rdkit_mol=rd_mol, conformer_coordinates=self.coordinates
        )
        Chem.AssignStereochemistryFrom3D(rd_mol)

        # now we should check that the stereo has not been broken
        for rd_atom in rd_mol.GetAtoms():
            index = rd_atom.GetIdx()
            qb_atom = self.atoms[index]
            if qb_atom.stereochemistry is not None:
                with _assert_wrapper(StereoChemistryError):
                    assert qb_atom.stereochemistry == rd_atom.GetProp(
                        "_CIPCode"
                    ), f"StereoChemistry incorrect expected {qb_atom.stereochemistry} got {rd_atom.GetProp('_CIPCode')} for atom {qb_atom}"

        for rd_bond in rd_mol.GetBonds():
            index = rd_bond.GetIdx()
            qb_bond = self.bonds[index]
            if qb_bond.stereochemistry is not None:
                rd_bond.SetStereo(qb_bond.rdkit_stereo)
            rd_stereo = rd_bond.GetStereo()
            if qb_bond.stereochemistry == "E":
                with _assert_wrapper(StereoChemistryError):
                    assert (
                        rd_stereo == Chem.BondStereo.STEREOE
                    ), f"StereoChemistry incorrect expected E got {rd_stereo}"
            elif qb_bond.stereochemistry == "Z":
                with _assert_wrapper(StereoChemistryError):
                    assert (
                        rd_stereo == Chem.BondStereo.STEREOZ
                    ), f"StereoChemistry incorrect expected Z got {rd_stereo}"

        return Chem.Mol(rd_mol)

    def get_smarts_matches(self, smirks: str) -> Optional[List[Tuple[int, ...]]]:
        """
        Get substructure matches for a mapped SMARTS pattern.

        Args:
            smirks:
                The mapped SMARTS pattern that should be used to query the molecule.

        Returns:
            `None` if there are no matches, else a list of tuples of atom indices which match the tagged atoms in
            the SMARTS pattern. These are returned in the same order.
        """
        matches = RDKit.get_smirks_matches(rdkit_mol=self.to_rdkit(), smirks=smirks)
        if not matches:
            return None
        return matches

    def add_qm_scan(self, scan_data: TorsionDriveData) -> None:
        """
        Save the torsion drive data into the ligand object.
        """
        if scan_data.__class__ != TorsionDriveData:
            raise MissingReferenceData(
                "The reference data must be in the form of the torsion drive data class."
            )
        else:
            if self.qm_scans is None:
                self.qm_scans = []
            self.qm_scans.append(scan_data)

    def openmm_coordinates(self) -> unit.Quantity:
        """
        Convert the coordinates to an openMM quantity.

        Build a single set of coordinates for the molecule that work in openMM.
        Note this must be a single conformer, if multiple are given only the first is used.

        Returns:
            A openMM quantity wrapped array of the coordinates in angstrom.
        """
        return unit.Quantity(self.coordinates, unit.angstroms)

    def fix_net_charge(self):
        """
        Ensure the total is exactly equal to the ideal net charge of the molecule.
        If net charge is not an integer value, MM simulations can (ex/im)plode.
        """

        decimal.setcontext(decimal.Context(prec=7))
        round_to = decimal.Decimal(10) ** -6

        for param in self.NonbondedForce:
            param.charge = param.charge.quantize(round_to)

        atom_charges = sum(param.charge for param in self.NonbondedForce)
        extra = self.charge - atom_charges

        if self.extra_sites is not None:
            for site in self.extra_sites:
                site.charge = site.charge.quantize(round_to)
                extra -= site.charge

        if extra:
            last_atom_index = self.n_atoms - 1
            self.NonbondedForce[(last_atom_index,)].charge += extra

class Ligand(Molecule):
    """
    The Ligand class seperats from protiens as we add fields to store QM calculations, such as the hessian and add more
    rdkit support methods.
    """

    hessian: Optional[Array[float]] = Field(
        None,
        description="The hessian matrix calculated for this molecule at the QM optimised geometry.",
    )
    qm_scans: Optional[List[TorsionDriveData]] = Field(
        None,
        description="The list of reference torsiondrive results which we can fit against.",
    )
    wbo: Optional[Array[float]] = Field(
        None, description="The WBO matrix calculated at the QM optimised geometry."
    )

    def __init__(
        self,
        atoms: List[Atom],
        bonds: Optional[List[Bond]] = None,
        coordinates: Optional[np.ndarray] = None,
        multiplicity: int = 1,
        name: str = "unk",
        routine: Optional[Set] = None,
        provenance: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        super(Ligand, self).__init__(
            atoms=atoms,
            bonds=bonds,
            coordinates=coordinates,
            multiplicity=multiplicity,
            name=name,
            routine=routine,
            provenance=provenance,
            **kwargs,
        )
        # make sure we have unique atom names
        self._validate_atom_names()

    @validator("hessian", "wbo", allow_reuse=True)
    def _reshape_matrix(cls, matrix: Optional[np.ndarray]) -> Optional[np.ndarray]:
        if matrix is not None:
            if len(matrix.shape) == 1:
                # the matrix is a flat list
                # so we need to make the matrix to be square
                length = int(np.sqrt(matrix.shape[0]))
                return matrix.reshape((length, length))
        return matrix

    @classmethod
    def from_rdkit(
        cls, rdkit_mol: Chem.Mol, name: Optional[str] = None, multiplicity: int = 1
    ) -> "Ligand":
        """
        Build an instance of a qubekit ligand directly from an rdkit molecule.

        Args:
            rdkit_mol:
                An instance of an rdkit.Chem.Mol from which the QUBEKit ligand should be built.
            name:
                The name that should be assigned to the molecule, this will overwrite any name already assigned.
            multiplicity:
                The multiplicity of the molecule, used in QM calculations.
        """
        if name is None:
            if rdkit_mol.HasProp("_Name"):
                name = rdkit_mol.GetProp("_Name")

        atoms = []
        bonds = []
        # Collect the atom names and bonds
        for rd_atom in rdkit_mol.GetAtoms():
            # make and atom
            qb_atom = Atom.from_rdkit(rd_atom=rd_atom)
            atoms.append(qb_atom)

        # now we need to make a list of bonds
        for rd_bond in rdkit_mol.GetBonds():
            qb_bond = Bond.from_rdkit(rd_bond=rd_bond)
            bonds.append(qb_bond)

        coords = rdkit_mol.GetConformer().GetPositions()
        bonds = bonds or None
        # method use to make the molecule
        routine = {"QUBEKit.ligand.from_rdkit"}
        return cls(
            atoms=atoms,
            bonds=bonds,
            coordinates=coords,
            multiplicity=multiplicity,
            name=name,
            routine=routine,
        )

    def has_ub_terms(self) -> bool:
        """Return `True` if the molecule has Ure-Bradly terms, as there are forces between non-bonded atoms."""
        for bond in self.BondForce:
            try:
                self.get_bond_between(*bond.atoms)
            except TopologyMismatch:
                return True
        return False

    def _to_ub_pdb(self, file_name: Optional[str] = None) -> None:
        """A privet method to write the molecule to a non-standard pdb file with connections for Urey-Bradly terms."""
        openmm_top = self.to_openmm_topology()
        PDBFile.writeFile(
            topology=openmm_top,
            positions=self.openmm_coordinates(),
            file=open(f"{file_name or self.name}.pdb", "w"),
        )

    @staticmethod
    def _check_file_name(file_name: str) -> None:
        """
        Make sure that if an unsupported file type is passed we can not make a molecule from it.
        """
        if ".xyz" in file_name:
            raise FileTypeError(
                "XYZ files can not be used to build ligands due to ambiguous bonding, "
                "please use pdb, mol, mol2 or smiles as input."
            )

    @classmethod
    def from_file(cls, file_name: str, multiplicity: int = 1) -> "Ligand":
        """
        Build a ligand from a supported input file.

        Args:
            file_name:
                The abs path to the file including the extension which determines how the file is read.
            multiplicity:
                The multiplicity of the molecule which is required for QM calculations.
        """
        cls._check_file_name(file_name=file_name)
        input_data = ReadInput.from_file(file_name=file_name)
        ligand = cls.from_rdkit(
            rdkit_mol=input_data.rdkit_mol,
            name=input_data.name,
            multiplicity=multiplicity,
        )
        # now edit the routine to include this call
        ligand.provenance["routine"].update(
            ["QUBEKit.ligand.from_file", os.path.abspath(file_name)]
        )
        return ligand

    @classmethod
    def from_smiles(
        cls, smiles_string: str, name: str, multiplicity: int = 1
    ) -> "Ligand":
        """
        Build the ligand molecule directly from a non mapped smiles string.

        Args:
            smiles_string:
                The smiles string from which a molecule instance should be made.
            name:
                The name that should be assigned to the molecule.
            multiplicity:
                The multiplicity of the molecule, important for QM calculations.
        """
        input_data = ReadInput.from_smiles(smiles=smiles_string, name=name)
        ligand = cls.from_rdkit(
            rdkit_mol=input_data.rdkit_mol, name=name, multiplicity=multiplicity
        )
        # now edit the routine to include this command
        ligand.provenance["routine"].update(
            ["QUBEKit.ligand.from_smiles", smiles_string]
        )
        return ligand

    def to_openmm_topology(self) -> Topology:
        """
        Convert the Molecule to a OpenMM topology representation.

        We assume we have a single molecule so a single chain is made with a single residue.
        Note this will not work with proteins as we will need to have distinct residues.

        Returns:
            An openMM topology object which can be used to construct a system.
        """
        topology = Topology()
        bond_types = {1: Single, 2: Double, 3: Triple}
        chain = topology.addChain()
        # create a molecule specific residue
        residue = topology.addResidue(name=self.name, chain=chain)
        # add atoms and keep track so we can add bonds
        top_atoms = []
        for atom in self.atoms:
            element = Element.getByAtomicNumber(atom.atomic_number)
            top_atom = topology.addAtom(
                name=atom.atom_name, element=element, residue=residue
            )
            top_atoms.append(top_atom)
        for bond in self.bonds:
            atom1 = top_atoms[bond.atom1_index]
            atom2 = top_atoms[bond.atom2_index]
            # work out the type
            if bond.aromatic:
                b_type = Aromatic
            else:
                b_type = bond_types[bond.bond_order]
            topology.addBond(
                atom1=atom1, atom2=atom2, type=b_type, order=bond.bond_order
            )
        # now check for Urey-Bradley terms
        for bond in self.BondForce:
            try:
                self.get_bond_between(*bond.atoms)
            except TopologyMismatch:
                # the bond is not in the bond list so add it as a u-b term
                atom1 = top_atoms[bond.atoms[0]]
                atom2 = top_atoms[bond.atoms[1]]
                # this is a fake bond used for U-B terms.
                topology.addBond(atom1=atom1, atom2=atom2, type=Single, order=1)

        return topology

    def to_smiles(
        self,
        isomeric: bool = True,
        explicit_hydrogens: bool = True,
        mapped: bool = False,
    ) -> str:
        """
        Create a canonical smiles representation for the molecule based on the input settings.

        Args:
            isomeric:
                If the smiles string should encode stereochemistry `True` or not `False`.
            explicit_hydrogens:
                If hydrogens should be explicitly encoded into the smiles string `True` or not `False`.
            mapped:
                If the smiles should encode the original atom ordering `True` or not `False` as this might be different
                from the canonical ordering.

        Returns:
            A smiles string which encodes the molecule with the desired settings.
        """
        return RDKit.get_smiles(
            rdkit_mol=self.to_rdkit(),
            isomeric=isomeric,
            explicit_hydrogens=explicit_hydrogens,
            mapped=mapped,
        )

    def generate_atom_names(self) -> None:
        """
        Generate a unique set of atom names for the molecule.
        """
        atom_names = {}
        for atom in self.atoms:
            symbol = atom.atomic_symbol
            if symbol not in atom_names:
                atom_names[symbol] = 1
            else:
                atom_names[symbol] += 1

            atom.atom_name = f"{symbol}{atom_names[symbol]}"

    def _validate_atom_names(self) -> None:
        """
        Check that the ligand has unique atom names if not generate a new set.
        """
        if not self.has_unique_atom_names:
            self.generate_atom_names()

    def to_qcschema(self, extras: Optional[Dict] = None) -> qcel.models.Molecule:
        """
        build a qcschema molecule from the ligand object, this is useful to interface with QCEngine and QCArchive.
        """
        import copy

        # make sure we have a conformer
        if self.coordinates == [] or self.coordinates is None:
            raise ConformerError(
                "The molecule must have a conformation to make a qcschema molecule."
            )
        coords = copy.deepcopy(self.coordinates)
        # input must be in bohr
        coords *= constants.ANGS_TO_BOHR
        # we do not store explicit bond order so guess at 1
        bonds = [
            (bond.atom1_index, bond.atom2_index, bond.bond_order) for bond in self.bonds
        ]
        mapped_smiles = self.to_smiles(
            isomeric=True, explicit_hydrogens=True, mapped=True
        )
        if extras is not None:
            extras["canonical_isomeric_explicit_hydrogen_mapped_smiles"] = mapped_smiles
        else:
            extras = {
                "canonical_isomeric_explicit_hydrogen_mapped_smiles": mapped_smiles
            }

        symbols = [atom.atomic_symbol for atom in self.atoms]
        schema_info = {
            "symbols": symbols,
            "geometry": coords,
            "connectivity": bonds,
            "molecular_charge": self.charge,
            "molecular_multiplicity": self.multiplicity,
            "extras": extras,
            "fix_com": True,
            "fix_orientation": True,
            "fix_symmetry": "c1",
        }
        return qcel.models.Molecule.from_data(schema_info, validate=True)

    def add_conformer(self, file_name: str) -> None:
        """
        Read the given input file extract  the conformers and save them to the ligand.
        TODO do we want to check that the connectivity is the same?
        """
        input_data = ReadInput.from_file(file_name=file_name)
        if input_data.coords is None:
            # get the coords from the rdkit molecule
            coords = input_data.rdkit_mol.GetConformer().GetPositions()
        else:
            if isinstance(input_data.coords, list):
                coords = input_data.coords[-1]
            else:
                coords = input_data.coords
        self.coordinates = coords


# -*- coding: utf-8 -*-
"""Data wrangling .ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1xZB07hyN9-KTMuFhnXZ1eBOaT2MyPdHI

# **Space X  Falcon 9 First Stage Landing Prediction**

## Lab 2: Data wrangling

Estimated time needed: **60** minutes

In this lab, we will perform some Exploratory Data Analysis (EDA) to find some patterns in the data and determine what would be the label for training supervised models.

In the data set, there are several different cases where the booster did not land successfully. Sometimes a landing was attempted but failed due to an accident; for example, True Ocean means the mission outcome was successfully  landed to a specific region of the ocean while False Ocean means the mission outcome was unsuccessfully landed to a specific region of the ocean. True RTLS means the mission outcome was successfully  landed to a ground pad False RTLS means the mission outcome was unsuccessfully landed to a ground pad.True ASDS means the mission outcome was successfully landed on  a drone ship False ASDS means the mission outcome was unsuccessfully landed on a drone ship.

In this lab we will mainly convert those outcomes into Training Labels with `1` means the booster successfully landed `0` means it was unsuccessful.

Falcon 9 first stage will land successfully

![](https://cf-courses-data.s3.us.cloud-object-storage.appdomain.cloud/IBMDeveloperSkillsNetwork-DS0701EN-SkillsNetwork/api/Images/landing\_1.gif)

Several examples of an unsuccessful landing are shown here:

![](https://cf-courses-data.s3.us.cloud-object-storage.appdomain.cloud/IBMDeveloperSkillsNetwork-DS0701EN-SkillsNetwork/api/Images/crash.gif)

## Objectives

Perform exploratory  Data Analysis and determine Training Labels

*   Exploratory Data Analysis
*   Determine Training Labels

***

## Import Libraries and Define Auxiliary Functions

We will import the following libraries.
"""

# Pandas is a software library written for the Python programming language for data manipulation and analysis.
import pandas as pd
#NumPy is a library for the Python programming language, adding support for large, multi-dimensional arrays and matrices, along with a large collection of high-level mathematical functions to operate on these arrays
import numpy as np

"""### Data Analysis

Load Space X dataset, from last section.
"""

df=pd.read_csv("https://cf-courses-data.s3.us.cloud-object-storage.appdomain.cloud/IBM-DS0321EN-SkillsNetwork/datasets/dataset_part_1.csv")
df.head(10)

"""Identify and calculate the percentage of the missing values in each attribute

"""

df.isnull().sum()/df.count()*100

"""Identify which columns are numerical and categorical:

"""

df.dtypes

"""### TASK 1: Calculate the number of launches on each site

The data contains several Space X  launch facilities: Cape Canaveral Space Launch Complex 40  VAFB SLC 4E  , Vandenberg Air Force Base Space Launch Complex 4E (SLC-4E), Kennedy Space Center Launch Complex 39A KSC LC 39A .The location of each Launch Is placed in the column LaunchSite

Next, let's see the number of launches for each site.

Use the method  value_counts() on the column LaunchSite to determine the number of launches  on each site:
"""

# Apply value_counts() on column LaunchSite
df["LaunchSite"].value_counts()

"""Each launch aims to an dedicated orbit, and here are some common orbit types:

*   LEO: Low Earth orbit (LEO)is an Earth-centred orbit with an altitude of 2,000 km (1,200 mi) or less (approximately one-third of the radius of Earth),\[1] or with at least 11.25 periods per day (an orbital period of 128 minutes or less) and an eccentricity less than 0.25.\[2] Most of the manmade objects in outer space are in LEO \[1].

*   VLEO: Very Low Earth Orbits (VLEO) can be defined as the orbits with a mean altitude below 450 km. Operating in these orbits can provide a number of benefits to Earth observation spacecraft as the spacecraft operates closer to the observation\[2].

*   GTO A geosynchronous orbit is a high Earth orbit that allows satellites to match Earth's rotation. Located at 22,236 miles (35,786 kilometers) above Earth's equator, this position is a valuable spot for monitoring weather, communications and surveillance. Because the satellite orbits at the same speed that the Earth is turning, the satellite seems to stay in place over a single longitude, though it may drift north to south,” NASA wrote on its Earth Observatory website \[3] .

*   SSO (or SO): It is a Sun-synchronous orbit  also called a heliosynchronous orbit is a nearly polar orbit around a planet, in which the satellite passes over any given point of the planet's surface at the same local mean solar time \[4] .

*   ES-L1 :At the Lagrange points the gravitational forces of the two large bodies cancel out in such a way that a small object placed in orbit there is in equilibrium relative to the center of mass of the large bodies. L1 is one such point between the sun and the earth \[5] .

*   HEO A highly elliptical orbit, is an elliptic orbit with high eccentricity, usually referring to one around Earth \[6].

*    ISS  A modular space station (habitable artificial satellite) in low Earth orbit. It is a multinational collaborative project between five participating space agencies: NASA (United States), Roscosmos (Russia), JAXA (Japan), ESA (Europe), and CSA (Canada) \[7]

*    MEO  Geocentric orbits ranging in altitude from 2,000 km (1,200 mi) to just below geosynchronous orbit at 35,786 kilometers (22,236 mi). Also known as an intermediate circular orbit. These are "most commonly at 20,200 kilometers (12,600 mi), or 20,650 kilometers (12,830 mi), with an orbital period of 12 hours  \[8]

*    HEO  Geocentric orbits above the altitude of geosynchronous orbit (35,786 km or 22,236 mi)  \[9]

*    GEO  It is a circular geosynchronous orbit 35,786 kilometres (22,236 miles) above Earth's equator and following the direction of Earth's rotation  \[10]

*    PO  It is one type of satellites in which a satellite passes above or nearly above both poles of the body being orbited (usually a planet such as the Earth  \[11]

some are shown in the following plot:

![](https://cf-courses-data.s3.us.cloud-object-storage.appdomain.cloud/IBMDeveloperSkillsNetwork-DS0701EN-SkillsNetwork/api/Images/Orbits.png)

### TASK 2: Calculate the number and occurrence of each orbit

Use the method  .value_counts() to determine the number and occurrence of each orbit in the  column Orbit
"""

# Apply value_counts on Orbit column
df["Orbit"].value_counts()

"""### TASK 3: Calculate the number and occurence of mission outcome per orbit type

Use the method  value_counts() to determine the number and occurrence of each orbit in the  column Outcome , then assign it to the variable landing_outcomes:
"""

# landing_outcomes = values on Outcome column
landing_outcomes = df["Outcome"].value_counts()
landing_outcomes

"""True Ocean means the mission outcome was successfully  landed to a specific region of the ocean while False Ocean means the mission outcome was unsuccessfully landed to a specific region of the ocean. True RTLS means the mission outcome was successfully  landed to a ground pad False RTLS means the mission outcome was unsuccessfully landed to a ground pad.True ASDS means the mission outcome was successfully  landed to a drone ship False ASDS means the mission outcome was unsuccessfully landed to a drone ship. None ASDS and None None these represent a failure to land.

"""

for i,outcome in enumerate(landing_outcomes.keys()):
    print(i,outcome)

"""We create a set of outcomes where the second stage did not land successfully:

"""

bad_outcomes=set(landing_outcomes.keys()[[1,3,5,6,7]])
bad_outcomes

"""### TASK 4: Create a landing outcome label from Outcome column

Using the Outcome,  create a list where the element is zero if the corresponding  row  in  Outcome is in the set bad_outcome; otherwise, it's one. Then assign it to the variable landing_class:
"""

# landing_class = 0 if bad_outcome
# landing_class = 1 otherwise

def onehot(item):
    if item in bad_outcomes:
        return 0
    else:
        return 1
landing_class = df["Outcome"].apply(onehot)
landing_class

"""This variable will represent the classification variable that represents the outcome of each launch. If the value is zero, the  first stage did not land successfully; one means  the first stage landed Successfully

"""

df['Class']=landing_class
df[['Class']].head(8)

df.head(5)

"""We can use the following line of code to determine  the success rate:

"""

df["Class"].mean()

"""We can now export it to a CSV for the next section,but to make the answers consistent, in the next lab we will provide data in a pre-selected date range.

df.to_csv("dataset_part\_2.csv", index=False)

## Authors

Joseph Santarcangelo has a PhD in Electrical Engineering, his research focused on using machine learning, signal processing, and computer vision to determine how videos impact human cognition. Joseph has been working for IBM since he completed his PhD.

Nayef Abou Tayoun is a Data Scientist at IBM and pursuing a Master of Management in Artificial intelligence degree at Queen's University.

## Change Log

| Date (YYYY-MM-DD) | Version | Changed By | Change Description      |
| ----------------- | ------- | ---------- | ----------------------- |
| 2020-09-20        | 1.0     | Joseph     | Modified Multiple Areas |
| 2020-11-04        | 1.1.    | Nayef      | updating the input data |
| 2021-05-026       | 1.1.    | Joseph     | updating the input data |

(づ｡◕‿‿◕｡)づ(づ｡◕‿‿◕｡)づ(づ｡◕‿‿◕｡)づ
"""

06 - Capstone Project/Week 1 Introduction/data_wrangling_.py


import os
import numpy as np
import cv2
import pickle
import torch
import imageio
from PIL import Image
import data.eval_utils as eval_utils
from data.util import RHD2Bighand_skeidx, root_idx, wrist_idx, norm_idx
from data.base_dataset import BaseDataset

class RHDDateset(BaseDataset):
    def __init__(self, path_name, mode, input_size = 64, output_size = 64, bbs_scale = 1.1, output_full = False, aug = True, sigma=3):
        super().__init__(path_name, 'rhd', mode, input_size, output_size, bbs_scale, output_full, aug)

        self.sigma = sigma
        with open(os.path.join(path_name + self.mode, 'anno_%s.pickle' % self.mode), 'rb') as fi:
            self.anno_all = pickle.load(fi)
            self.num_samples = len(self.anno_all)

    def __len__(self):
        return int(self.num_samples)

    def __getitem__(self, idx):
        if (idx == 20500 or idx == 28140): idx = 0

        anno = self.anno_all[idx]
        image = imageio.imread(os.path.join(self.path_name, self.mode, 'color', '%.5d.png' % idx))
        mask = imageio.imread(os.path.join(self.path_name, self.mode, 'mask', '%.5d.png' % idx))
        uv = anno['uv_vis'][:, :2]
        kp_visible = anno['uv_vis'][:, 2] == 1
        pose3d = anno['xyz']
        camera = anno['K']

        hand_image, hand_mask, new_uv, crop_center, crop_size, hand_side, uv_vis, \
        relative_depth, pose3d_root, bone_length, pose3d, pose_uv_all, pose3d_normed\
            = self.segment_hand(image, mask, uv, kp_visible, pose3d)
        image = np.array(image, dtype=np.float)

        rot_mat_inv = np.eye(3, 2)
        rot_mat = np.eye(2, 3)
        hand_image_aug, hand_mask_aug, new_uv_aug, pose3d_normed_aug \
            = hand_image.copy(), hand_mask.copy(),new_uv.copy(), pose3d_normed.copy()

        if self.aug:
            hand_image_aug, hand_mask_aug, new_uv_aug, rot_mat, rot_mat_inv, pose3d_normed_aug \
                = self.augmentation(hand_image, hand_mask, new_uv, pose3d_normed)

        hand_image_tensor = Image.fromarray(hand_image_aug)
        hand_image_tensor = self.image_trans(hand_image_tensor).float()

        target = {}
        target['img_crop'] = hand_image_tensor
        target['uv_crop'] = torch.from_numpy(new_uv_aug).view(21,2).float()
        target['vis21'] = torch.from_numpy(uv_vis).view(21).float()
        target['relative_depth'] = torch.from_numpy(relative_depth).view(21,1).float()
        target['mask'] = torch.as_tensor(hand_mask_aug).view(1,self.output_size,self.output_size).float()
        target['confidence'] = torch.ones(1).view(1).float()

        hm_res = 64
        hm = np.zeros(
            (21, hm_res, hm_res),
            dtype='float32'
        )  # (CHW)
        hm_veil = np.ones(21, dtype='float32')
        uv_crop = target['uv_crop'].numpy()
        for i in range(21):
            # kp = (
            #         (uv_crop[i] / crop_size) * hm_res
            # ).astype(np.int32)  # kp uv: [0~256] -> [0~64]
            hm[i], aval = eval_utils.gen_heatmap(hm[i], uv_crop[i], self.sigma)
            hm_veil[i] *= aval
        target['hm'] = hm
        target['hm_veil'] = hm_veil

        if self.output_full:
            target['camera_type'] = 'perspective'  # orthographic
            target['pose3d_normed'] = torch.from_numpy(pose3d_normed_aug).float()
            target['rot_mat_inv'] = torch.from_numpy(rot_mat_inv).float()
            target['rot_mat'] = torch.from_numpy(rot_mat).float()
            target['pose3d'] = torch.from_numpy(pose3d).float()
            target['pose_uv_all'] = torch.from_numpy(pose_uv_all).float()
            target['crop_center'], target['crop_size'], target['hand_side'], target['pose3d_root'], target['bone_length'], target['camera'] = \
                torch.from_numpy(crop_center).float(),torch.as_tensor(crop_size).float(),\
                torch.as_tensor(hand_side).bool(),torch.as_tensor(pose3d_root).float(),torch.as_tensor(bone_length).float(),\
                torch.from_numpy(camera).float()

        return target

    def segment_hand(self, image, mask, kp_coord_uv, kp_visible, pose3d):
        cond_l = np.logical_and(mask > 1, mask < 18)
        cond_r = mask > 17
        num_px_left_hand = np.sum(cond_l)
        num_px_right_hand = np.sum(cond_r)
        hand_side = num_px_left_hand > num_px_right_hand

        ### crop image
        if hand_side:
            pose_uv_all = kp_coord_uv[:21, :]
            uv_vis = kp_visible[:21]
            pose3d = pose3d[:21]
        else:
            pose_uv_all = kp_coord_uv[-21:, :]
            uv_vis = kp_visible[-21:]
            pose3d = pose3d[-21:]

        if hand_side:
            maskSingleHand = cond_l.astype(np.float)
        else:
            maskSingleHand = cond_r.astype(np.float)

        ### RHD2bighand idx
        pose3d = pose3d[RHD2Bighand_skeidx,:]
        pose_uv_all = pose_uv_all[RHD2Bighand_skeidx,:]
        uv_vis = uv_vis[RHD2Bighand_skeidx]

        crop_center = pose_uv_all[root_idx, :]
        crop_center = np.reshape(crop_center, 2)

        pose_uv_vis = pose_uv_all[uv_vis, :]
        crop_size = np.max(np.absolute(pose_uv_vis - crop_center))
        crop_size = np.minimum(np.maximum(crop_size, 25.0), 200.0)
        crop_size = np.ceil(crop_size) * self.bbs_scale

        # print crop_size
        image_crop = self.imcrop(image, crop_center, crop_size)
        mask_crop = self.imcrop(maskSingleHand, crop_center, crop_size)

        # resize image to the input size of hourglass
        height, width, _ = image.shape
        refine_size = self.output_size
        image_crop = cv2.resize(image_crop, (self.input_size, self.input_size), interpolation=cv2.INTER_NEAREST)
        mask_crop = cv2.resize(mask_crop, (refine_size, refine_size), interpolation=cv2.INTER_NEAREST)

        crop_scale = refine_size / (crop_size * 2)
        new_uv = self.creat_refine_uv(pose_uv_all, crop_center, refine_size, crop_scale)

        ### get relative depth
        pose3d_root = pose3d[root_idx, :]  # this is the root coord
        pose3d_rel = pose3d - pose3d_root  # relative coords in metric coords
        bone_length = np.sqrt(np.sum(np.square(pose3d_rel[root_idx, :] - pose3d_rel[norm_idx, :])))
        pose3d_normed = pose3d_rel / bone_length
        relative_depth = pose3d_normed[:,2]

        if hand_side:  # transfer left hands to right hand
            image_crop = cv2.flip(image_crop, 1)
            mask_crop = cv2.flip(mask_crop,1)
            new_uv[:, 0] = refine_size - new_uv[:, 0]
            pose3d_normed[:, 0] = -pose3d_normed[:, 0]

        return image_crop, mask_crop, new_uv, crop_center, crop_size, hand_side, uv_vis, \
               relative_depth, pose3d_root, bone_length, pose3d, pose_uv_all, pose3d_normed

if __name__ == '__main__':
    import os,torch
    os.chdir('/home/yangl/MultiTaskHand/')

    train_dataset = RHDDateset('../dataset/RHD_published_v2/', 'training', input_size=256, output_full=True)
    Dataloader = torch.utils.data.DataLoader(train_dataset,
                                            batch_size=2,
                                            shuffle=True,
                                            drop_last=True,
                                            num_workers=0)

    from util import crop2xyz

    iter = iter(Dataloader)

    for i in range(1000):
        print(i)
        _, _, label_rgb, target = next(iter)

        uv_crop = target['new_uv']
        norm_depth_pose = target['relative_depth']
        uv_original, xyz = crop2xyz(uv_crop, norm_depth_pose, target, 64)
        print((uv_original - target['pose_uv_all']).max())
        print((xyz - target['pose3d']).max())

        # hand_image = hand_image_tensor.permute(1, 2, 0).numpy()
        # hand_image = (hand_image + 1) / 2
        # #
        # # pose3d = target['pose3d'].numpy()
        # # if target['hand_side']:
        # #     pose3d[:, 0] = -pose3d[:, 0]
        #
        # # mask = target['mask'].squeeze().cpu().numpy()
        # #
        # # print(target['new_uv'])
        # #
        # # print(target['relative_depth'])
        #
        # import matplotlib.pyplot as plt
        # from mpl_toolkits.mplot3d import Axes3D
        # from vis.skeleton import plot_pose3d
        # fig = plt.figure()
        # ax1 = fig.add_subplot(111)
        # ax1.imshow(hand_image)
        # ax1.axis('off')
        # # plt.show()
        # plt.savefig('./rhd/{}.png'.format(i),bbox_inches='tight')
        # plt.close()

        # fig = plt.figure()
        # ax1 = fig.add_subplot(111, projection='3d')
        # plot_pose3d(pose3d, '.-', ax1, azim=-90.0, elev=180.0)
        # ax1.axis('off')
        # plt.show()
        # # plt.savefig('rhd.png',bbox_inches='tight')
        # plt.close()

    # import matplotlib.pyplot as plt
    # from vis.skeleton import mask_cmap
    # fig = plt.figure()
    # ax1 = fig.add_subplot(111)
    # ax1.imshow(mask,cmap=mask_cmap)
    # ax1.axis('off')
    # plt.savefig('rhd_mask.png',bbox_inches='tight')
    # plt.close()