# Sidewalk Extrapolation

import os
import json
from shapely.geometry import Point, LineString
import pandas as pd
import geopandas as gpd
from geopandas import GeoSeries, GeoDataFrame

import matplotlib.pyplot as plt
import matplotlib
import fiona

# The two statemens below are used mainly to set up a plotting
# default style that's better than the default from matplotlib
import seaborn as sns
plt.style.use('bmh')

json_config = open("config.json").read()
c = json.loads(json_config)

# REMEMBER TO USE PYTHON 3:
# I found it useful to use a virtualenv for this
# to create: virtualenv -p /usr/bin/python3 py3env
# to activate: source py3env/bin/activate

def main():
	streets_ntwk = gpd.read_file(os.path.join(c['data_path'], c['sidewalk_network_shapefile']))
	lr_stations = gpd.read_file(os.path.join(c['data_path'], c['station_shapefile']))

	station_buffers = {}
	for station in c['target_stations']:
		station_buffers[station] = buffer_station(streets_ntwk, lr_stations, station, c['buffer_distance_feet'])

	#plot_target_stations(swk_network, lr_stations, station_buffers)

	plot_ref = lr_stations.plot(marker='*', color='red', markersize=5)
	for station_key in station_buffers.keys():
		streets = station_buffers[station_key]
		streets.plot(ax=plot_ref, color='grey')
		swks = generate_swks(streets)
		plot_swks(swks, plot_ref)
	plt.show()

# returns geoseries of generated sidewalks within the passes buffers
def generate_swks(streets):
	sidewalks = []
	streets = filter_sidewalks(streets)
	for idx, row in streets.iterrows():
		street = row[c['geometry']]
		col_name = c['st_type_code']
		print(col_name)
		st_type_intermediate = row[col_name]
		st_type = c['st_type'][st_type_intermediate]
		offset = c['st_type_offset'][st_type]
		print(offset)
		if row[c['swk_left']] == c['swk_present']:
			sidewalks.append(street.parallel_offset(offset, 'left', resolution=16, join_style=1, mitre_limit=40.0))
		if row[c['swk_right']] == c['swk_present']:
			sidewalks.append(street.parallel_offset(offset, 'right', resolution=16, join_style=1, mitre_limit=40.0))
	return sidewalks

# hacky way to plot swks to handle size
def plot_swks(swks, plot_ref):
	series_builder = []
	print(type(swks))
	idx = 0
	for sidewalk in swks:
		series_builder.append(sidewalk)
		if idx % 100 == 0:
			GeoSeries(series_builder).plot(ax=plot_ref, color='blue')
			series_builder = []
		idx += 1

def plot_target_stations(swk_network, lr_stations, station_buffers):
	# retrieve plot_ref so future plot calls can be put on same figure
	plot_ref = lr_stations.plot(marker='*', color='red', markersize=5)
	# swk_network is too large to plot with pyplot
	# swk_network.plot(ax=plot_ref, color='grey', linewidth=0.3)
	# swk_ntwrk.plot(ax=plot_ref, color='grey')
	for station_key in station_buffers.keys():
		station = station_buffers[station_key]
		filter_sidewalks(station).plot(ax=plot_ref, color='blue')
	# labels for stations
	for idx, row in lr_stations.iterrows():
		point = row[c['geometry']]
		coord = (point.x, point.y)
		plt.annotate(s=row['name'], xy=coord, fontsize=8)
	plt.show()

# returns only streets with a sidewalk on at least one side
def filter_sidewalks(df_swks):
	only_swks = df_swks.query(c['swk_left'] + ' == ' + str(c['swk_present']) + ' or ' + c['swk_right'] + ' == ' + str(c['swk_present']))
	return only_swks

# returns dataframe with only swks surrounding station at the specified distance
# Note: the two geodata frames must have the same crs
def buffer_station(swk_ntwrk, stations, station_name, distance):
	station = stations[stations.name == station_name]
	buff = station.buffer(distance).iloc[0]
	intersection = swk_ntwrk.loc[swk_ntwrk.intersects(buff)]
	return intersection

# takes column name of line geometry to sum and column name of sidewalk left and right presence
# Note: sidewalk presence is specified by 1, 0 means absence.
# returns the sum of of all sidewalks in meters
def sum_sidewalks(gdf, sum_row_name):
	sidewalk_sum = 0.0
	for index, row in gdf.iterrows():
		geo = row[sum_row_name]
		road_length = geo.length
		if row[c['swk_left']] == c['swk_present']:
			sidewalk_sum += road_length
		if row[c['swk_right']] == c['swk_present']:
			sidewalk_sum += road_length
	return sidewalk_sum

if __name__ == "__main__":
    main()


#!/usr/bin/env python3
# coding: utf-8
"""
Created on Thu Apr 21 2022
@author: Simon Delarue <sdelarue@enst.fr>
"""
from typing import Union

import numpy as np
from scipy import sparse

from sknetwork.gnn.activation import get_activation_function
from sknetwork.gnn.base_layer import BaseLayer
from sknetwork.gnn.utils import has_self_loops, add_self_loops
from sknetwork.linalg import diag_pinv

class GCNConv(BaseLayer):
    """Graph convolutional operator.

    :math:`H^{\prime}=\sigma(D^{-1/2}\hat{A}D^{-1/2}HW + b)`,

    where :math:`\hat{A} = A + I` denotes the adjacency matrix with inserted self-loops and
    :math:`D` its diagonal degree matrix. :math:`W` and :math:`b` are trainable parameters and
    :math:`\\sigma` is the activation function.

    Parameters
    ----------
    out_channels: int
        Size of each output sample.
    activation: str (default = ``'Relu'``)
        Activation function.
        Can be either:

        * ``'relu'``, the rectified linear unit function, returns f(x) = max(0, x)
        * ``'sigmoid'``, the logistic sigmoid function, returns f(x) = 1 / (1 + exp(-x)).
        * ``'softmax'``, the softmax function, returns f(x) = exp(x) / sum(exp(x))
    use_bias: bool (default = `True`)
        If ``True``, add a bias vector.
    normalization: str (default = ``'Both'``)
        Normalization of the adjacency matrix for message passing.
        Can be either:

        * ``'left'``, left normalization by the vector of degrees
        * ``'right'``, right normalization by the vector of degrees
        * ``'both'``,  symmetric normalization by the square root of degrees
    self_loops: bool (default = `True`)
        If ``True``, add self-loops to each node in the graph.

    Attributes
    ----------
    weight: np.ndarray,
        Trainable weight matrix.
    bias: np.ndarray
        Bias vector.
    update: np.ndarray
        Embedding of the nodes before the activation function.
    embedding: np.ndarray
        Embedding of the nodes after convolution layer.

    References
    ----------
    Kipf, T., & Welling, M. (2017).
    `Semi-supervised Classification with Graph Convolutional Networks.
    <https://arxiv.org/pdf/1609.02907.pdf>`_
    5th International Conference on Learning Representations.

    He, K. & Zhang, X. & Ren, S. & Sun, J. (2015).
    `Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification.
    <https://arxiv.org/pdf/1502.01852.pdf>`_
    Proceedings of the IEEE International Conference on Computer Vision (ICCV).
    """
    def __init__(self, out_channels: int, activation: str = 'Relu', use_bias: bool = True,
                 normalization: str = 'Both', self_loops: bool = True):
        super(GCNConv, self).__init__(out_channels, activation, use_bias, normalization, self_loops)

    def forward(self, adjacency: Union[sparse.csr_matrix, np.ndarray],
                features: Union[sparse.csr_matrix, np.ndarray]) -> np.ndarray:
        """Compute graph convolution.

        Parameters
        ----------
        adjacency
            Adjacency matrix of the graph.
        features : sparse.csr_matrix, np.ndarray
            Input feature of shape :math:`(n, d)` with :math:`n` the number of nodes in the graph and :math:`d`
            the size of feature space.

        Returns
        -------
        embedding: p.ndarray
            Node embedding.
        """
        if not self.weights_initialized:
            self._initialize_weights(features.shape[1])

        n_row, n_col = adjacency.shape

        if self.self_loops:
            if not has_self_loops(adjacency):
                adjacency = add_self_loops(adjacency)

        weights = adjacency.dot(np.ones(n_col))
        if self.normalization == 'left':
            d_inv = diag_pinv(weights)
            adjacency = d_inv.dot(adjacency)
        elif self.normalization == 'right':
            d_inv = diag_pinv(weights)
            adjacency = adjacency.dot(d_inv)
        elif self.normalization == 'both':
            d_inv = diag_pinv(np.sqrt(weights))
            adjacency = d_inv.dot(adjacency).dot(d_inv)

        msg = adjacency.dot(features)
        update = msg.dot(self.weight)

        if self.use_bias:
            update += self.bias

        activation_function = get_activation_function(self.activation)
        embedding = activation_function(update)

        # Keep track of results for backprop
        self.embedding = embedding
        self.update = update

        return embedding


''' Tools for prepare images for machine learning '''
import os
import cv2
import numpy as np
import matplotlib.pyplot as plt

def slice_image(img, x_size, y_size):
    images = []
    y = 0
    while y < img.shape[0]:
        y_end = min(y + y_size, img.shape[0])
        x = 0
        while x < img.shape[1]:
            x_end = min(x + x_size, img.shape[1])
            images.append(img[x:x_end, y:y_end])
            x = x_end
        y = y_end
    return images

def batch_write(file_name, images):
    filename, ext = os.path.splitext(file_name)
    for image, count in zip(images, range(len(images))):
        name = "{}.{:010}{}".format(filename, count, ext)
        cv2.imwrite(name, image)

def slice_images(in_path, out_path, width, height):
    os.makedirs(out_path, exist_ok=True)
    for file in sorted(os.listdir(in_path)):
        if file.endswith(".tif"):
            image = cv2.imread(os.path.join(in_path, file), cv2.IMREAD_GRAYSCALE)
            images = slice_image(image, width, height)
            batch_write(os.path.join(out_path, file), images)

def remove_empty(empty_file_path, relative_file_path, minimum_ratio = 0.1):
    for file in os.listdir(empty_file_path):
        if file.endswith(".tif"):
            empty_file = os.path.join(empty_file_path, file)
            other_file = os.path.join(relative_file_path, file)
            image = cv2.imread(empty_file, cv2.IMREAD_GRAYSCALE)
            hist = np.histogram(image.flatten(), 256, [0, 256])[0]
            if hist[255]/(hist[0]+hist[255]) < minimum_ratio:
                os.remove(empty_file)
                os.remove(other_file)

def merge_images(images,filename):  #merge
    images = np.array(images)
    n_img = len(images)
    size_img = images[0].shape
    n_img_0 = int(np.sqrt(n_img))
    if n_img_0 != np.floor(n_img_0):
        print('Error!, amount of images is not square!')
    # np.reshape(images,(None,10*128))
    image = images.reshape(n_img_0,n_img_0,size_img[0],size_img[1]).transpose(1,0,2,3).swapaxes(1,2).reshape(size_img[0]*n_img_0,-1)
    # print(image[0:30,0:30],'\n',type(image),image.shape, n_img, size_img)

    fig88 = plt.figure(88)
    plt.imshow(image, cmap='gray')
    fig88.suptitle('Merged image')
    # plt.savefig(filename)
    # plt.show()
    image[np.nonzero(image)] = 255
    image = np.uint8(image)  # to gray
    cv2.imwrite(filename, image)

    clay_pixel = np.count_nonzero(image==255)
    total_pixel = np.prod(image.shape)
    print('Predicted image saved. \nThe percentage of Clay/organic is: ',100*clay_pixel/total_pixel,'%\n')


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import pickle
from functools import partial
from time import sleep

from pyproxy_client import PyProxyClient, new_simple_connection

from .future import Future, future_id

class PyProxySession:
    """
    PyProxySession is a class for establishing the mainstream
    with the PyProxy Server.
    """
    def __init__(self, addr="localhost:9000"):
        self._addr = addr

        # block - waitint for client conenction
        # raise exception if we fail
        self._client_conn = new_simple_connection(addr)

    @property
    def session_id(self):
        """
        session_id, as sent in the server-hello message
        """
        return self._client_conn.session_id

    def __enter__(self):
        """
        """
        return self.connect()

    def connect(self):
        """
        connect will spawn a background thread
        we return a RemoteProcess object
        """
        client = PyProxyClient(self._client_conn)
        return RemoteProcess(client)

    def __exit__(self, exc_typ, exc_val, trcb):
        self._client.disconnect()

class RemoteProcess:
    """
    RemoteProcess is our object wrapping a single "remote python process",
    we can send code to execute and interact with stdin/stdout and stderr.
    """
    def __init__(self, client):
        self._client = client

    def output(self, future=None):
        """
        output retrieves stdout and stderr line from the remote process
        """

        break_now = True
        if future:
            break_now = False

        while True:
            # out is a tuple
            # first arg is 1 or 2 (stdout or stderr)
            # second arg is a string for one pipe line
            out = self._client.next_output()
            if out:
                yield out
                continue

            if break_now:
                break

            if future and future.is_done():
                # break on next loop iteration
                break_now = True

            if not break_now:
                # sleep for 100ms then poll again
                sleep(0.1)

    def stdin(self, lines):
        """
        stdin will send lines to stdin on the remote process.
        Currently we raise NotImplementedError
        """

        raise NotImplementedError("stdin not implemented")

    def eval(self, code, locs=None, globs=None):
        """
        eval will execute a code object on the remote process
        code may be a str or a code object

        if code_object is a string then we send to the server as a String
        if code_object is a code_object then we attempt to pickle
        """

        id = future_id()

        locs = locs or {}
        globs = globs or {}

        locs = pickle.dumps(locs)
        globs = pickle.dumps(globs)

        if isinstance(code, str):
            inner_fut = self._client.eval_str(id, code, locs, globs)
        else:
            raise TypeError("code must be a string")

        return Future(id, inner_fut)

    def disconnect(self):
        self._client.disconnect()


import matplotlib.pyplot as plt
from numpy import arange

def plot_loss_acc(losses, accuracies, epochs):

    train_loss, val_loss = losses
    train_acc, val_acc = accuracies
    epochs_range = range(1, epochs +1)
    plt.figure(1)
    plt.plot(epochs_range, train_loss, 'r-', label='train')
    plt.plot(epochs_range, val_loss, 'b-', label='val')
    plt.xlabel('Epoch')
    plt.ylabel('Loss value')
    plt.title('Loss')
    plt.savefig('loss.jpg')

    plt.figure(2)
    plt.plot(epochs_range, train_acc, 'r-', label='train')
    plt.plot(epochs_range, val_acc, 'b-', label='val')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.title('Training/Validation Accuracy')
    plt.savefig('accuracy.jpg')


# In memory of Do Deeee

from datetime import date
import os
from fastapi import FastAPI, Depends, Query, Body
from pydantic import BaseModel
from typing import Optional
from app.models import PostSchema, UserLoginSchema, UserSchema
from app.auth.auth_handler import sign_JWT
from app.auth.auth_bearer import JWTBearer

app = FastAPI(
    title="BayApp Developer API",
    description="This is an API for the BayApp. Contact dev at khizer.khan98@gmail.com for more information.",
    version="Alpha Pre Release 0.1"
)

posts = [
    {
        "id" : 1,
        "title" : "Pancakes",
        "content" : "Lorem Ipsum"
    }
]

users = []

@app.get("/")
async def read_items():
    return {'msg': 'Hello'}

@app.get('/posts')
async def get_posts():
    return {'data' : posts}

@app.get('/posts/{id}')
async def get_single_post(id:int):
    if id > len(posts):
        return {'error' : 'No post with supplied ID'}
    else:
        for post in posts:
            if post["id"] == id:
                return {'data' : post
                }

@app.post('/add_post', dependencies=[Depends(JWTBearer())])
async def add_post(post:PostSchema):
    post.id = len(posts) + 1
    posts.append(post.dict())

    return {'data' : 'post added'}

@app.post('/user/signup')
async def create_user(user : UserSchema):
    users.append(user)
    return sign_JWT(user.email)

@app.post('/user/login')
async def user_login(user : UserLoginSchema):
    if check_user(user):
        return sign_JWT(user.email)
    return {"error" : "Wrong credentials"}

def check_user(data : UserLoginSchema):
    for user in users:
        if user.email == data.email and user.password == data.password:
            return True
    return False


from django.urls import path
from django.urls.resolvers import _PATH_PARAMETER_COMPONENT_RE
from . import views
from aplicaciones.citas.views import crearCitas, citas, medicos, crearMedicos, pacientes, crearPacientes,eliminarCitas,eliminarMedicos,eliminarPacientes, editarCitas, editarMedicos, editarPacientes

urlpatterns = [
    # crear rutas de citas
    path ('', views.home), path('', views.citas),
    path ('citas/', citas,name = 'citas'),
    path ('citas_crear/', crearCitas,name = 'citas_crear'),
    path ('citas_eliminar/<int:cit_id>', eliminarCitas,name = 'citas_eliminar'),
    path ('citas_editar/<int:cit_id>', editarCitas,name ='citas_editar'),

    # crear rutas de medicos
    path ('medicos/', medicos,name = 'medicos'),
    path ('medicos_crear/', crearMedicos,name = 'medicos_crear'),
    path ('medicos_eliminar/<int:med_id>', eliminarMedicos,name = 'medicos_eliminar'),
    path ('medicos_editar/<int:med_id>', editarMedicos,name ='medicos_editar'),

    # crear rutas pacientes
    path ('pacientes/', pacientes,name = 'pacientes'),
    path ('pacientes_crear/', crearPacientes,name = 'pacientes_crear'),
    path ('pacientes_eliminar/<int:pac_id>', eliminarPacientes,name = 'pacientes_eliminar'),
    path ('pacientes_editar/<int:med_id>', editarPacientes,name ='pacientes_editar'),

]


from __future__ import annotations
import os
from typing import NamedTuple, List
import multiprocessing as mp
from pathlib import Path
import warnings
import numpy as np
import pandas as pd
from tqdm.autonotebook import tqdm
from lobsterpy.featurize.core import (
    FeaturizeLobsterpy,
    FeaturizeCharges,
    FeaturizeCOXX,
)

warnings.filterwarnings("ignore")

class BatchSummaryFeaturizer:
    """
    Batch Featurizer sets that generates summary features from lobster data.

    Args:
        path_to_lobster_calcs: path to root directory consisting of all lobster calc
        path_to_jsons: path to root directory consisting of all lobster lightweight jsons
        feature_type: set the feature type for moment features.
        Possible options are "bonding", "antibonding" or "overall"
        charge_type : set charge type used for computing ionicity. Possible options are
        "mulliken", "loewdin or "both"
        bonds: "all_bonds" or "cation_anion_bonds"
        include_cobi_data : bool stating to include COBICAR.lobster features
        include_coop_data: bool stating to include COOPCAR.lobster features
        e_range : range of energy relative to fermi for which moment features needs to be computed
        n_jobs : parallel processes to run

    Attributes:
        get_df: A pandas dataframe with summary features
    """

    def __init__(
        self,
        path_to_lobster_calcs: str,
        path_to_jsons: str | None = None,
        feature_type: str = "antibonding",
        charge_type: str = "both",
        bonds: str = "all",
        include_cobi_data: bool = False,
        include_coop_data: bool = False,
        e_range: List[float] = [-5.0, 0.0],
        n_jobs: int = 4,
    ):
        self.path_to_lobster_calcs = path_to_lobster_calcs
        self.path_to_jsons = path_to_jsons
        self.feature_type = feature_type
        self.charge_type = charge_type
        self.bonds = bonds
        self.include_cobi_data = include_cobi_data
        self.include_coop_data = include_coop_data
        self.e_range = e_range
        self.n_jobs = n_jobs

    def _featurizelobsterpy(self, file_name_or_path) -> pd.DataFrame:
        """
        Wrapper method to featurize Lobsterpy condensed bonding analysis data by loading lightweight json
        if json file exists or invokes lobsterpy.analzye.Analysis module

        Returns:
            A pandas dataframe with ICOHP stats like mean, min, max of relevant bonds and
            madelung energies

        """
        if Path(file_name_or_path).is_file():
            featurize_lobsterpy = FeaturizeLobsterpy(
                path_to_json=file_name_or_path,
                bonds=self.bonds,
            )

        else:
            featurize_lobsterpy = FeaturizeLobsterpy(
                path_to_lobster_calc=file_name_or_path,
                bonds=self.bonds,
            )

        df = featurize_lobsterpy.get_df()

        return df

    def _featurizecoxx(self, path_to_lobster_calc) -> pd.DataFrame:
        """
        Wrapper method to featurize COHP/COBI/COOPCAR data that uses FeaturizeCOXX under the hood

        Returns:
            A pandas dataframe with COHP summary stats data mainly weighted ICOHP/ICOOP/ICOBI,
            Effective interaction number and moment features (center, width, skewness and kurtosis)

        """
        dir_name = Path(path_to_lobster_calc)

        req_files = {
            "structure_path": "POSCAR",
            "coxxcar_path": "COHPCAR.lobster",
            "icoxxlist_path": "ICOHPLIST.lobster",
        }
        for file, default_value in req_files.items():
            file_path = dir_name / default_value
            req_files[file] = file_path  # type: ignore
            if not file_path.exists():
                gz_file_path = file_path.with_name(file_path.name + ".gz")
                if gz_file_path.exists():
                    req_files[file] = gz_file_path  # type: ignore

        coxxcar_path = req_files.get("coxxcar_path")
        structure_path = req_files.get("structure_path")
        icoxxlist_path = req_files.get("icoxxlist_path")

        if (
            coxxcar_path.exists()  # type: ignore
            and structure_path.exists()  # type: ignore
            and icoxxlist_path.exists()  # type: ignore
        ):
            coxx = FeaturizeCOXX(
                path_to_coxxcar=str(coxxcar_path),
                path_to_icoxxlist=str(icoxxlist_path),
                path_to_structure=str(structure_path),
                feature_type=self.feature_type,
                e_range=self.e_range,
            )

            df_cohp = coxx.get_summarized_coxx_df()
        else:
            raise Exception(
                "COHPCAR.lobster or POSCAR or ICOHPLIST.lobster file "
                "not found in {}".format(dir_name.name)
            )

        if self.include_cobi_data:
            req_files = {
                "coxxcar_path": "COBICAR.lobster",
                "icoxxlist_path": "ICOBILIST.lobster",
            }
            for file, default_value in req_files.items():
                file_path = dir_name / default_value
                req_files[file] = file_path  # type: ignore
                if not file_path.exists():
                    gz_file_path = file_path.with_name(file_path.name + ".gz")
                    if gz_file_path.exists():
                        req_files[file] = gz_file_path  # type: ignore

            coxxcar_path = req_files.get("coxxcar_path")
            icoxxlist_path = req_files.get("icoxxlist_path")

            if coxxcar_path.exists() and icoxxlist_path.exists():  # type: ignore
                coxx = FeaturizeCOXX(
                    path_to_coxxcar=str(coxxcar_path),
                    path_to_icoxxlist=str(icoxxlist_path),
                    path_to_structure=str(structure_path),
                    feature_type=self.feature_type,
                    e_range=self.e_range,
                    are_cobis=True,
                )

                df_cobi = coxx.get_summarized_coxx_df()

            else:
                raise Exception(
                    "COBICAR.lobster or ICOBILIST.lobster file "
                    "not found in {}".format(dir_name.name)
                )

        if self.include_coop_data:
            req_files = {
                "coxxcar_path": "COOPCAR.lobster",
                "icoxxlist_path": "ICOOPLIST.lobster",
            }
            for file, default_value in req_files.items():
                file_path = dir_name / default_value
                req_files[file] = file_path  # type: ignore
                if not file_path.exists():
                    gz_file_path = file_path.with_name(file_path.name + ".gz")
                    if gz_file_path.exists():
                        req_files[file] = gz_file_path  # type: ignore

            coxxcar_path = req_files.get("coxxcar_path")
            icoxxlist_path = req_files.get("icoxxlist_path")

            if coxxcar_path.exists() and icoxxlist_path.exists():  # type: ignore
                coxx = FeaturizeCOXX(
                    path_to_coxxcar=str(coxxcar_path),
                    path_to_icoxxlist=str(icoxxlist_path),
                    path_to_structure=str(structure_path),
                    feature_type=self.feature_type,
                    e_range=self.e_range,
                    are_coops=True,
                )

                df_coop = coxx.get_summarized_coxx_df()

            else:
                raise Exception(
                    "COOPCAR.lobster or ICOOPLIST.lobster file "
                    "not found in {}".format(dir_name.name)
                )

        if self.include_cobi_data and self.include_coop_data:
            df = pd.concat([df_cohp, df_cobi, df_coop], axis=1)
        elif self.include_cobi_data and not self.include_coop_data:
            df = pd.concat([df_cohp, df_cobi], axis=1)
        elif not self.include_cobi_data and self.include_coop_data:
            df = pd.concat([df_cohp, df_coop], axis=1)
        else:
            df = df_cohp

        return df

    def _featurizecharges(self, path_to_lobster_calc) -> pd.DataFrame:
        """
        Wrapper method to featurize CHARGE.lobster.gz data that uses FeaturizeCharges under the hood

        Returns:
            A pandas dataframe with computed ionicity for the structure

        """
        dir_name = Path(path_to_lobster_calc)

        req_files = {
            "charge_path": "CHARGE.lobster",
            "structure_path": "POSCAR",
        }
        for file, default_value in req_files.items():
            file_path = dir_name / default_value
            req_files[file] = file_path  # type: ignore
            if not file_path.exists():
                gz_file_path = file_path.with_name(file_path.name + ".gz")
                if gz_file_path.exists():
                    req_files[file] = gz_file_path  # type: ignore

        charge_path = req_files.get("charge_path")
        structure_path = req_files.get("structure_path")

        if charge_path.exists() and structure_path.exists():  # type: ignore
            if self.charge_type == "mulliken":
                charge_mull = FeaturizeCharges(
                    path_to_charge=str(charge_path),
                    path_to_structure=str(structure_path),
                    charge_type="mulliken",
                )
                df = charge_mull.get_df()
            elif self.charge_type == "loewdin":
                charge_loew = FeaturizeCharges(
                    path_to_charge=str(charge_path),
                    path_to_structure=str(structure_path),
                    charge_type="loewdin",
                )
                df = charge_loew.get_df()
            elif self.charge_type == "both":
                charge_mull = FeaturizeCharges(
                    path_to_charge=str(charge_path),
                    path_to_structure=str(structure_path),
                    charge_type="mulliken",
                )
                df_mull = charge_mull.get_df()

                charge_loew = FeaturizeCharges(
                    path_to_charge=str(charge_path),
                    path_to_structure=str(structure_path),
                    charge_type="loewdin",
                )
                df_loew = charge_loew.get_df()

                df = pd.concat([df_mull, df_loew], axis=1)

            return df
        else:
            raise Exception(
                "CHARGE.lobster or POSCAR not found in {}".format(dir_name.name)
            )

    def get_df(self) -> pd.DataFrame:
        """
        This method will return a pandas dataframe with summary features extracted from LOBSTER files
        as columns. Uses multiprocessing to speed up the process.

        Returns:
            Returns a pandas dataframe

        """
        if self.path_to_jsons:
            file_name_or_path = [
                os.path.join(self.path_to_jsons, f)
                for f in os.listdir(self.path_to_jsons)
                if not f.startswith("t")
                and not f.startswith(".")
                and not os.path.isdir(f)
            ]

        elif self.path_to_lobster_calcs and not self.path_to_jsons:
            file_name_or_path = [
                os.path.join(self.path_to_lobster_calcs, f)
                for f in os.listdir(self.path_to_lobster_calcs)
                if not f.startswith("t")
                and not f.startswith(".")
                and os.path.isdir(os.path.join(self.path_to_lobster_calcs, f))
            ]

        with mp.Pool(processes=self.n_jobs, maxtasksperchild=1) as pool:
            results = tqdm(
                pool.imap_unordered(
                    self._featurizelobsterpy, file_name_or_path, chunksize=1
                ),
                total=len(file_name_or_path),
                desc="Generating LobsterPy summary stats",
            )
            pool.close()
            pool.join()
            row = []
            for result in results:
                row.append(result)

        df_lobsterpy = pd.concat(row)
        df_lobsterpy.sort_index(inplace=True)

        paths = [
            os.path.join(self.path_to_lobster_calcs, f)
            for f in os.listdir(self.path_to_lobster_calcs)
            if not f.startswith("t")
            and not f.startswith(".")
            and os.path.isdir(os.path.join(self.path_to_lobster_calcs, f))
        ]

        with mp.Pool(processes=self.n_jobs, maxtasksperchild=1) as pool:
            results = tqdm(
                pool.imap_unordered(self._featurizecoxx, paths, chunksize=1),
                total=len(paths),
                desc="Generating COHP/COOP/COBI summary stats",
            )
            pool.close()
            pool.join()
            row = []
            for result in results:
                row.append(result)

        df_coxx = pd.concat(row)
        df_coxx.sort_index(inplace=True)

        with mp.Pool(processes=self.n_jobs, maxtasksperchild=1) as pool:
            results = tqdm(
                pool.imap_unordered(self._featurizecharges, paths, chunksize=1),
                total=len(paths),
                desc="Generating charge based features",
            )
            pool.close()
            pool.join()
            row = []
            for result in results:
                row.append(result)

        df_charges = pd.concat(row)
        df_charges.sort_index(inplace=True)

        df = pd.concat([df_lobsterpy, df_coxx, df_charges], axis=1)

        return df

class BatchCoxxFingerprint:
    """
    BatchFeaturizer that generates COHP/COOP/COBI fingerprints and
    Tanimoto index similarity matrix from fingerprint objects.

    Args:
        path_to_lobster_calcs: path to root directory consisting of all lobster calc
        feature_type: set the feature type for moment features.
        Possible options are "bonding", "antibonding" or "overall"
        label_list: bond labels list for which fingerprints needs to be generated.
        tanimoto : bool to state to compute tanimoto index betweeen fingerprint objects
        normalize: bool to state to normalize the fingerprint data
        n_bins: sets number for bins for fingerprint objects
        e_range : range of energy relative to fermi for which moment features needs to be computed
        n_jobs : number of parallel processes to run
        fingerprint_for: Possible options are 'cohp/cobi/coop'.
        Based on this fingerprints will be computed for COHPCAR/COOBICAR/COOPCAR.lobster files

    Attributes:
        fingerprint_df: A pandas dataframe with fingerprint objects
        get_similarity_matrix_df: A symmetric pandas dataframe consisting of
        similarity index (tanimoto/normalized dot product/dot product)
        computed between all pairs of compunds
    """

    def __init__(
        self,
        path_to_lobster_calcs: str,
        feature_type: str = "overall",
        label_list: List[str] | None = None,
        tanimoto: bool = True,
        normalize: bool = True,
        spin_type: str = "summed",
        n_bins: int = 56,
        e_range: List[float] = [-15.0, 0.0],
        n_jobs=4,
        fingerprint_for: str = "cohp",
    ):
        self.path_to_lobster_calcs = path_to_lobster_calcs
        self.feature_type = feature_type
        self.tanimoto = tanimoto
        self.normalize = normalize
        self.label_list = label_list
        self.spin_type = spin_type
        self.n_bins = n_bins
        self.e_range = e_range
        self.n_jobs = n_jobs
        self.fingerprint_for = fingerprint_for

        self.fingerprint_df = self._get_fingerprints_df()

    def get_similarity_matrix_df(self) -> pd.DataFrame:
        """
        This function will compute pairwise similarity index for each fingerprint object in input dataframe

        Returns:
             A Pandas dataframe
        """
        matrix = np.full(
            (self.fingerprint_df.shape[0], self.fingerprint_df.shape[0]), np.nan
        )
        for i, (row, col) in enumerate(self.fingerprint_df.iterrows()):
            for j, (row1, col1) in enumerate(self.fingerprint_df.iterrows()):
                if self.tanimoto:
                    simi = self._get_fp_similarity(
                        col["COXX_FP"],
                        col1["COXX_FP"],
                        tanimoto=self.tanimoto,
                        normalize=False,
                    )
                else:
                    simi = self._get_fp_similarity(
                        col["COXX_FP"],
                        col1["COXX_FP"],
                        tanimoto=False,
                        normalize=True,
                    )
                matrix[i][j] = simi

        df = pd.DataFrame(
            matrix,
            index=list(self.fingerprint_df.index),
            columns=list(self.fingerprint_df.index),
        )

        return df

    @staticmethod
    def _fp_to_dict(fp) -> dict:
        """
        Converts a fingerprint into a dictionary

        Args:
            fp: The fingerprint to be converted into a dictionary

        Returns:
            dict: A dict of the fingerprint Keys=type, Values=np.ndarray(energies, cohp)
        """
        fp_dict = {}
        fp_dict[fp[2]] = np.array([fp[0], fp[1]], dtype="object").T

        return fp_dict

    @staticmethod
    def _get_fp_similarity(
        fp1: NamedTuple,
        fp2: NamedTuple,
        col: int = 1,
        pt: int | str = "All",
        normalize: bool = False,
        tanimoto: bool = True,
    ) -> float:
        """
        Calculates the similarity index (dot product) of two fingerprints

        Args:
            fp1 (NamedTuple): The 1st dos fingerprint object
            fp2 (NamedTuple): The 2nd dos fingerprint object
            col (int): The item in the fingerprints (0:energies,1: coxxs) to take the dot product of (default is 1)
            pt (int or str) : The index of the point that the dot product is to be taken (default is All)
            normalize (bool): If True normalize the scalar product to 1 (default is False)
            tanimoto (bool): If True will compute Tanimoto index (default is False)

        Raises:
            ValueError: If both tanimoto and normalize are set to True.

        Returns:
            Similarity index (float): The value of dot product

        """
        fp1_dict = (
            BatchCoxxFingerprint._fp_to_dict(fp1) if not isinstance(fp1, dict) else fp1
        )

        fp2_dict = (
            BatchCoxxFingerprint._fp_to_dict(fp2) if not isinstance(fp2, dict) else fp2
        )

        if pt == "All":
            vec1 = np.array([pt[col] for pt in fp1_dict.values()]).flatten()
            vec2 = np.array([pt[col] for pt in fp2_dict.values()]).flatten()
        else:
            vec1 = fp1_dict[fp1[2][pt]][col]
            vec2 = fp2_dict[fp2[2][pt]][col]

        if not normalize and tanimoto:
            rescale = (
                np.linalg.norm(vec1) ** 2
                + np.linalg.norm(vec2) ** 2
                - np.dot(vec1, vec2)
            )
            return np.dot(vec1, vec2) / rescale

        elif not tanimoto and normalize:
            rescale = np.linalg.norm(vec1) * np.linalg.norm(vec2)
            return np.dot(vec1, vec2) / rescale

        elif not tanimoto and not normalize:
            rescale = 1.0
            return np.dot(vec1, vec2) / rescale

        else:
            raise ValueError(
                "Cannot compute similarity index. Please set either normalize=True or tanimoto=True or both to False."
            )

    def _fingerprint_df(self, path_to_lobster_calc) -> pd.DataFrame:
        """
        Wrapper method to get fingerprint object dataframe using FeaturizeCOXX.get_coxx_fingerprint_df method.
        Also helps switching the data used for fingerprint generation

        Returns:
            A pandas dataframe with COXX fingerprint object

        """
        dir_name = Path(path_to_lobster_calc)

        if self.fingerprint_for.upper() == "COBI":
            req_files = {
                "coxxcar_path": "COBICAR.lobster",
                "icoxxlist_path": "ICOBILIST.lobster",
            }
            for file, default_value in req_files.items():
                file_path = dir_name / default_value
                req_files[file] = file_path  # type: ignore
                if not file_path.exists():
                    gz_file_path = file_path.with_name(file_path.name + ".gz")
                    if gz_file_path.exists():
                        req_files[file] = gz_file_path  # type: ignore

            coxxcar_path = req_files.get("coxxcar_path")
            icoxxlist_path = req_files.get("icoxxlist_path")
            are_cobis = True
            are_coops = False

        elif self.fingerprint_for.upper() == "COOP":
            req_files = {
                "coxxcar_path": "COOPCAR.lobster",
                "icoxxlist_path": "ICOOPLIST.lobster",
            }
            for file, default_value in req_files.items():
                file_path = dir_name / default_value
                req_files[file] = file_path  # type: ignore
                if not file_path.exists():
                    gz_file_path = file_path.with_name(file_path.name + ".gz")
                    if gz_file_path.exists():
                        req_files[file] = gz_file_path  # type: ignore

            coxxcar_path = req_files.get("coxxcar_path")
            icoxxlist_path = req_files.get("icoxxlist_path")
            are_cobis = False
            are_coops = True

        else:
            req_files = {
                "coxxcar_path": "COHPCAR.lobster",
                "icoxxlist_path": "ICOHPLIST.lobster",
            }
            for file, default_value in req_files.items():
                file_path = dir_name / default_value
                req_files[file] = file_path  # type: ignore
                if not file_path.exists():
                    gz_file_path = file_path.with_name(file_path.name + ".gz")
                    if gz_file_path.exists():
                        req_files[file] = gz_file_path  # type: ignore

            coxxcar_path = req_files.get("coxxcar_path")
            icoxxlist_path = req_files.get("icoxxlist_path")
            are_cobis = False
            are_coops = False

        structure_path = dir_name / "POSCAR"
        if not structure_path.exists():
            gz_file_path = structure_path.with_name(structure_path.name + ".gz")
            if gz_file_path.exists():
                structure_path = gz_file_path

        coxx = FeaturizeCOXX(
            path_to_coxxcar=str(coxxcar_path),
            path_to_icoxxlist=str(icoxxlist_path),
            path_to_structure=str(structure_path),
            feature_type=self.feature_type,
            e_range=self.e_range,
            are_coops=are_coops,
            are_cobis=are_cobis,
        )

        df_fp = coxx.get_coxx_fingerprint_df(
            spin_type=self.spin_type,
            n_bins=self.n_bins,
            normalize=self.normalize,
            label_list=self.label_list,
        )

        return df_fp

    def _get_fingerprints_df(self) -> pd.DataFrame:
        """
        Batch wrapper method to get fingerprint objects dataframe using
        BatchCoxxFingerprint._fingerprint_df method.

        Returns:
            A pandas dataframe with COXX fingerprint objects

        """
        paths = [
            os.path.join(self.path_to_lobster_calcs, f)
            for f in os.listdir(self.path_to_lobster_calcs)
            if not f.startswith("t")
            and not f.startswith(".")
            and os.path.isdir(os.path.join(self.path_to_lobster_calcs, f))
        ]

        with mp.Pool(processes=self.n_jobs, maxtasksperchild=1) as pool:
            results = tqdm(
                pool.imap_unordered(self._fingerprint_df, paths, chunksize=1),
                total=len(paths),
                desc="Generating {} fingerprints".format(self.fingerprint_for.upper()),
            )
            pool.close()
            pool.join()
            row = []
            for result in results:
                row.append(result)

        df = pd.concat(row)
        df.sort_index(inplace=True)

        return df


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

path = 'https://s3-api.us-geo.objectstorage.softlayer.net/cf-courses-data/CognitiveClass/DA0101EN/automobileEDA.csv'
df = pd.read_csv(path)
df.head()
from sklearn.linear_model import LinearRegression
lm = LinearRegression()
lm

X = df[['highway-mpg']]
Y = df[['price']]
lm.fit(X,Y)

Yhat = lm.predict(X)
Yhat[0:5]
lm.intercept_
lm.coef_

X =df[['engine-size']]
Y = df[['price']]
lm1 = lm.fit(df[['engine-size']],df[['price']])
lm1

lm1.intercept_
lm1.coef_

Yhat = lm1.predict(df[['engine-size']])
Yhat[0:10]

Yhat = -7963.34-166.86*X
Yhat[0:10]

Z = df[['horsepower','curb-weight','engine-size','highway-mpg']]
lm.fit(Z,df['price'])
lm.fit(Z,df[['price']])

lm.intercept_
lm.coef_
# Price = -15678.742628061467 + 52.65851272 x horsepower + 4.69878948 x curb-weight + 81.95906216 x engine-size + 33.58258185 x highway-mpg

Z2 = df[['normalized-losses','highway-mpg']]
lm.fit(Z2, df[['price']])
lm.intercept_
lm.coef_
# price = 38201.31 + 1.50*normalized-losses -820.45*highway-mpg

import seaborn as sns

width = 12
height = 10
plt.figure(figsize=(width,height))
sns.regplot(x='highway-mpg',y='price',data=df)
plt.ylim(0,)

plt.figure(figsize=(width,height))
sns.regplot(x='peak-rpm',y='price',data=df)
plt.ylim(0,)
# The same thing
plt.figure(figsize=(width,height))
sns.regplot(df['peak-rpm'],df['price'])
plt.ylim(0,)

df2 =df[['peak-rpm','highway-mpg','price']]
df2.head(20)
df2.corr()

df[['peak-rpm','highway-mpg','price']].corr()

width =12
height =10
plt.figure(figsize=(width,height))
sns.residplot(df['highway-mpg'],df['price'])
# sns.residplot(x='highway-mpg',y='price', data=df)
plt.show()

Yhat =lm.predict(Z)
plt.figure(figsize=(width,height))
ax1 = sns.distplot(df['price'], hist=False, color="r", label="Actual Value")
sns.distplot(Yhat, hist=False, color="b", label="Fitted Values" , ax=ax1)

sns.distplot(df['price'], hist = False, color='r', label='Actual Value')
sns.distplot(Yhat, hist=False, color='b', label='Fitted Values')   # This 2 lines of code have the same output as above, why that rather than this way

plt.title('Acutal va Fitted Value for Price')
plt.xlabel('Price (in dollars)')
plt.ylabel('Proportion of Cars')

plt.show()
plt.close()

def PlotPolly(model, independent_variable, dependent_variable, Name):
    x_new = np.linspace(15,55,100)
    y_new = model(x_new)

    plt.plot(independent_variable, dependent_variable, '.',x_new, y_new, '-')
    plt.title('Polynomial Fit With Matplotlib for Price ~ Length')
    ax= plt.gca()
    ax.set_facecolor((0.898,0.898,0.898))
    fig = plt.gcf()
    plt.xlabel(Name)
    plt.ylabel('Price of Cars')
    plt.ylim(0,)

    plt.show()
    plt.close()

x = df['highway-mpg']
y = df['price']

f = np.polyfit(x,y,3)
p = np.poly1d(f)
print(p)
PlotPolly(p, x, y, 'highway-mpg')
np.polyfit(x, y, 3)

f2 = np.polyfit(x, y, 11)
p2 = np.poly1d(f2)
print(p2)

PlotPolly(p2, x, y, 'highway-mpg')

from sklearn.preprocessing import PolynomialFeatures
pr = PolynomialFeatures(degree =2)
pr

Z_pr = pr.fit_transform(Z)
Z.shape
Z_pr.shape

from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

Input = [('scale',StandardScaler()), ('polynomial', PolynomialFeatures(include_bias=False)), ('model', LinearRegression())]
pipe = Pipeline(Input)
pipe.fit(Z,y)
ypipe = pipe.predict(Z)
ypipe[0:4]

Input2 =[('scale', StandardScaler()),('model', LinearRegression())]
pipe2 = Pipeline(Input2)
pipe2.fit(Z, y)
ypipe = pipe.predict(Z)
ypipe[0:10:2]
ypipe[0:10]

lm.fit(X, Y)
print('The R-square =', lm.score(X,Y))
Yhat = lm.predict(X)
print('The output of the first predicted value is:', Yhat[0:4])
from sklearn.metrics import mean_squared_error
mse = mean_squared_error(df['price'], Yhat)
print('The mean square error of price and predicted value is:',mse)

lm.fit(Z,df['price'])
print('The R-square =', lm.score(Z, df['price']))

Y_predict_multifit = lm.predict(Z)
print('The output of the first predicted value is:', Y_predict_multifit[0:4])
mse2 = mean_squared_error(df['price'], Y_predict_multifit)
print('The mean square error of price and predicted value is:', mse2)

from sklearn.metrics import r2_score
r_squared = r2_score(y, p(x))
# X = df[['engine-size']]
# Y = df[['price']]
# x = df['engine-size']
# y = df['price']
print('The R-square value is:', r_squared)

mean_squared_error(df['price'],p(x))

new_input = np.arange(1,100,1).reshape(-1,1)
lm.fit(X,Y)
lm
yhat = lm.predict(new_input)
yhat[0:5]
plt.plot(new_input, yhat)
plt.show()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""     Simulation functions -- Version 4.2
Last edit:  2022/09/24
Author(s):  Geysen, Steven (SG)
Notes:      - Functions used for the simulation of the task used by
                Marzecova et al. (2019). Both structure and models.
                * Models
                    - Rescorla-Wagner (Daphne)
                    - Rescorla-Wagner - Pearce-Hall hybrid (Hugo)
                    - Meta learner (Michelle)
                    - Win-stay-lose-shift (Wilhelm)
                    - Random (Renee)
                * (negaitve) log likelihood
                * Negative Spearman correlation
                * Accuracy
                * var_bin_switch
            - Release notes:
                * Changed reward calculations
                * Target in sim_experiment changed
                * Valitdity corrected
                * Negative correlation with absolute values

To do:      - Meta learner

Comments:   SG: Simulations return pandas.DataFrame. Model functions return
                originale data with model selection appended.

Sources:    https://www.frontiersin.org/articles/10.3389/fpsyg.2018.00612/full
            https://docs.sympy.org/latest/modules/stats.html
            https://github.com/Kingsford-Group/ribodisomepipeline/blob/master/scripts/exGaussian.py
            https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.exponnorm.html
            https://hannekedenouden.ruhosting.nl/RLtutorial/html/SoftMax.html
            https://lindeloev.shinyapps.io/shiny-rt/
"""

#%% ~~ Imports ~~ %%#

import numpy as np
import pandas as pd

import fns.assisting_functions as af

from scipy import stats

#%% ~~ Structure ~~ %%#
#######################

def sim_experiment(simnr=1, ntrials=640, nswitch=7):
    """
    Simulate the eperimental structure of Marzecova et al. (2019)

    Parameters
    ----------
    simnr : int, optional
        Simulation number. Used for the randomisation of the reward
            probabilities.
        The default is 1.
    ntrials : int, optional
        Number of trials.
        The default is 640.
    nswitch : int, optional
        Number of switches in cue validity during the experiment.
        The default is 7.

    Returns
    -------
    data : pandas.DataFrame
        Contains:
            0. 'id' - participants id
            1. 'trial' - trial # of the task
            2. 'relCueCol' - colour of the relevant cue (0: white / 1: black)
            3. 'relCue' - direction of the relevant cue (0: left / 1: right)
            4. 'irrelCue'- direction of the irrelevant cue (0: left / 1: right)
            5. 'targetLoc' - location of the target (0: left / 1: right)
            6. 'validity' - trial validity
            7. 'gammaBlock' - the validity level within the block
    """

    # Variables
    #----------
    # DataFrame
    column_list = [
        'id', 'trial', 'relCueCol', 'relCue', 'irrelCue', 'targetLoc',
        'validity', 'gammaBlock'
        ]
    dataDict = {keyi: [] for keyi in column_list}
    dataDict['id'] = simnr

    # Task
    n_stim = 2
    ## Colour of relevant cue
    relCueCol = np.random.randint(n_stim)
    ## Probabilities of reward for each stim
    prob = np.full(n_stim, 0.5)
    ## SG: For simulations with an even number is the initial probability
        ## of the relevant cue 0.7
    if simnr % 2 == 0:
        prob[relCueCol] = 0.7
    else:
        prob[relCueCol] = 0.85
    ## Trials where probability is switched
    switch = np.cumsum(np.random.randint(40, 120, size = nswitch))

    # Trial loop
    #-----------
    for triali in range(ntrials):
        dataDict['trial'].append(triali)
        # Switch probabilities
        if switch.shape[0] != 0:
            if triali == switch[0]:
                relCueCol = 1 - relCueCol

                if simnr % 2 == 0:
                    ##SG: For simulations with an even number is in the
                        # second half of trials the probability
                        # of the relevant cue 0.85.
                    if triali >= ntrials / 2:
                        prob[relCueCol] = 0.85
                    else:
                        prob[relCueCol] = 0.7
                else:
                    if triali >= ntrials / 2:
                        prob[relCueCol] = 0.7
                    else:
                        prob[relCueCol] = 0.85
                prob[1 - relCueCol] = 0.5
                ## Remove first item in switch
                switch = np.delete(switch, 0)
        dataDict['relCueCol'].append(relCueCol)
        dataDict['gammaBlock'].append(prob[relCueCol])

        # Stimuli
        stim = np.random.choice(n_stim, n_stim, p =  [0.5, 0.5])

        relcue = stim[relCueCol]
        irrelcue = stim[1 - relCueCol]
        dataDict['relCue'].append(relcue)
        dataDict['irrelCue'].append(irrelcue)

        ##SG: Target is 'randomly' selected between 0 and 1 with the
            # probability of the relevant cue.
        target = np.random.choice(
            [relcue, 1 - relcue], p = [prob[relCueCol], 1 -  prob[relCueCol]]
            )
        dataDict['targetLoc'].append(target)
        ##SG: Valid trials are coded as 0, invalid as 1. For consistancy with
            # behavioural data. Now normal
        dataDict['validity'].append(int(relcue == target))

    data = pd.DataFrame(dataDict, columns=column_list)

    return data

#%% ~~ Blocks ~~ %%#
####################

def sim_rt(rpe):
    """
    Simulate response times
    Response times, sampled from ExGaussian distribution. Sigma and mu are
    taken from exgauss fit of the original data. Tau is the reward prediction
    error of the selected cue, used as proxy for difficulty (absolute values to
    avoid errors caused by negative values).

    Parameters
    ----------
    rpe : int
        Reward prediction error of the selected cue.

    Returns
    -------
    RT : float
        Simulated response time (in secondes).
    """

    if rpe == 0:
        RT = 0.3
    else:
        try:
            ##SG: K = tau / sigma, loc = mu, scale = sigma.
            RT = stats.exponnorm.rvs(K = abs(rpe) / 0.02635,
                                     loc = 0.3009, scale = 0.02635)
        except:
            RT = np.nan
            print('Failed rt sampling')
        if RT == float('inf'):
            RT = 1.7

    return RT

def sim_models():
    """
    Simulation dictionary
    Contains the functions of the different simulation models.
    """

    return {'RW': simRW_1c,
            'RW2': simRW_2c,
            'H': simHybrid_1c,
            'M': simMeta_1c,
            'H2': simHybrid_2c,
            'W': simWSLS,
            'R': simRandom}

#%% ~~ Models ~~ %%#
####################

#%% ~~ Models (1 cue) ~~ %%#
#--------------------------#

# ~~ Rescorla - Wagner ~~ #
def simRW_1c(parameters, data, asm='soft'):
    """
    Daphne the delta learner
    The Rescorla - Wagner learning model with SoftMax. The policy calculations
    are part of the function (and not a separate function). Only the cue
    estimate of the selected cue is updated.

    Parameters
    ----------
    parameters : tuple, list, array
        First parameter is the learning rate of the model (0 <= alpha <= 1).
        Second parameter is the constant of the action selection method
            (0 < beta)..
    data : pandas.DataFrame
        Dataframe containing the structure of the experiment.
    asm : string, optional
        The action selection method, policy. The default is SoftMax.

    Returns
    -------
    simData : pandas.DataFrame
        Contains columns of 'data' and simulated behaviour of Daphne:
            0. 'selCue_RW' - selected cue by the Rescorla-Wagner model
            1. 'prob_RW' - probability to select cue 0
            2. 'rt_RW' - response time
            3. 'reward_RW' - reward based on cue selected by Rescorla-Wagner,
                also the validity of the selected cue
            4. 'RPE_RW' - reward prediction error
            5. 'Men_PE_RW' - Prediction error calculated as in
                Mengotti et al. (2017)
            6. 'Qest_0_RW' - estimated value of cue 0
            7. 'Qest_1_RW' - estimated value of cue 1
    """

    # Variables
    # ---------
    # Dataframe
    var_list = [
        'selCue_RW', 'prob_RW', 'rt_RW', 'reward_RW',
        'RPE_RW', 'Men_PE_RW', 'Qest_0_RW', 'Qest_1_RW'
        ]
    simDict = {vari:[] for vari in var_list}

    ## Number of trials
    n_trials = data.shape[0]
    ## Number of cues
    N_CUES = 2
    # Model
    ## Estimated value of cue
    Q_est = np.full((n_trials, N_CUES), 1/N_CUES)
    ## Outcomes
    outcomes = np.full((2, ), np.nan)

    # Policy
    ## Selected cue
    selcue = np.nan

    # Parameters
    alpha = parameters[0]
    beta, bias = 0, 0
    if len(parameters) > 1:
        beta = parameters[1]
    if len(parameters) > 2:
        bias = parameters[2]

    # Trial loop
    for triali, trial in data.iterrows():
        # Policy
        # ------
        ## Random for first trial
        if triali == 0:
            selcue = np.random.randint(N_CUES)
            probcue = 0.5
        else:
            selcue, probcue = af.policy(asm, Q_est[triali - 1, :],
                                        beta, bias)
        simDict['selCue_RW'].append(selcue)
        simDict['prob_RW'].append(probcue)

        # Reward
        # ------
        ## Based on validity
        ##AM: If cue==target reward = 1, if cue!=target reward = 0
        """
        AM: The model chooses the relevant cue color (1 - white, 2 - black),
        but you need to also know the direction of the cue
        (1 - left,  2 - right). The direction of the cue is in the column
        relCue or irrelCue, but the relCue or irrelCue refer to the gorund
        thruth and not to what the model chooses (selCue). RelCue and irrelCue
        don’t give you the information about which color is the relevant Cue
        currently, this information is in the relCueCol (1-white 2 -black). So,
        it would be easier to recode the data from the perspective of whiteCue
        and blackCue: if white cue is relevant, you would put the value for
        relCue direction in the column for white cue, while you would put the
        value for direction of the irelCue for the black cue.
        """
        outcomes[int(trial.relCueCol)] = trial.relCue == trial.targetLoc
        outcomes[int(1 - trial.relCueCol)] = trial.irrelCue == trial.targetLoc
        reward = int(outcomes[selcue])
        simDict['reward_RW'].append(reward)

        # Update rule (RW)
        # ----------------
        if triali == 0:
            # Reward prediction error
            rpe = reward - Q_est[triali, selcue]
        else:
            # Reward prediction error
            rpe = reward - Q_est[triali - 1, selcue]
            # Cue estimates
            ## Repeat cue estimates of previous trial
            Q_est[triali, :] = Q_est[triali - 1, :]
        ## Update cue estimate of selected stimulus in current trial
        Q_est[triali, selcue] = Q_est[triali, selcue] + alpha * rpe
        simDict['RPE_RW'].append(rpe)
        simDict['Men_PE_RW'].append(
            1 - reward + Q_est[triali - 1, selcue] * (2 * reward - 1)
            )
        for qi, q_est in enumerate(Q_est[triali, :]):
            simDict[f'Qest_{qi}_RW'].append(q_est)

        # Response time
        simDict['rt_RW'].append(sim_rt(rpe))

    return af.save_data(simDict, data, var_list)

# ~~ RW-PH Hybrid ~~ #
def simHybrid_1c(parameters, data, salpha=0.01, asm='soft'):
    """
    Hugo the hybrid learner
    The RW-PH hybrid learning model with SoftMax. The policy calculations
    are part of the function (and not a separate function). Only the cue
    estimate of the selected cue is updated.

    Parameters
    ----------
    parameters : tuple, list, array
        First parameter is eta (0 <= eta <= 1).
        Second parameter is the constant of the action selection method
            (0 < beta).
    data : pandas.DataFrame
        Dataframe containing the structure of the experiment.
    salpha : float, optional
        Start of alpha. The default is 0.01.
    asm : string, optional
        The action selection method, policy. The default is SoftMax.

    Returns
    -------
    simData : pandas.DataFrame
        Contains columns of 'data' and simulated behaviour of Hugo:
            0. 'selCue_H' - selected cue by the hybrid model
            1. 'prob_H' - probability to select cue 0
            2. 'rt_H' - response time
            3. 'reward_H' - reward based on cue selected by hybrid model
            4. 'alpha_H' - learning rate
            5. 'RPE_H' - reward prediction error
            6. 'Qest_0_H' - estimated value of cue 0
            7. 'Qest_1_H' - estimated value of cue 1
    """

    # Variables
    # ---------
    # Dataframe
    var_list = [
        'selCue_H', 'prob_H', 'rt_H', 'reward_H', 'alpha_H',
        'RPE_H', 'Qest_0_H', 'Qest_1_H'
        ]
    simDict = {vari:[] for vari in var_list}

    ## Number of trials
    n_trials = data.shape[0]
    ## Number of cues
    N_CUES = 2

    # Model
    ## Alpha
    alpha = np.full((n_trials, N_CUES), salpha)
    ## Estimated value of cue
    Q_est = np.full((n_trials, N_CUES), 1/N_CUES)
    ## Outcomes
    outcomes = np.full((2, ), np.nan)

    # Policy
    ## Selected cue
    selcue = np.nan

    # Parameters
    eta = parameters[0]
    beta, bias = 0, 0
    if len(parameters) > 1:
        beta = parameters[1]
    if len(parameters) > 2:
        bias = parameters[2]

    # Trial loop
    for triali, trial in data.iterrows():
        # Policy
        # ------
        ## Random for first trial
        if triali == 0:
            selcue = np.random.randint(N_CUES)
            probcue = 0.5
        else:
            selcue, probcue = af.policy(asm, Q_est[triali - 1, :],
                                        beta, bias)
        simDict['selCue_H'].append(selcue)
        simDict['prob_H'].append(probcue)

        # Reward
        # ------
        ## Based on validity
        ##AM: If cue==target reward = 1, if cue!=target reward = 0
        outcomes[int(trial.relCueCol)] = trial.relCue == trial.targetLoc
        outcomes[int(1 - trial.relCueCol)] = trial.irrelCue == trial.targetLoc
        reward = int(outcomes[selcue])
        simDict['reward_H'].append(reward)

        # Hybrid
        # ------
        if triali == 0:
            # Reward prediction error
            rpe = reward - Q_est[triali, selcue]
        else:
            # Reward prediction error
            rpe = reward - Q_est[triali - 1, selcue]
            # Update values of selected stimulus in current trial
            alpha[triali, :] = alpha[triali - 1, :]
            Q_est[triali, :] = Q_est[triali - 1, :]
        # Alpha (PH)
        alpha[triali, selcue] = eta * np.abs(rpe) + \
            (1 - eta) * alpha[triali, selcue]
        # Cue estimates
        Q_est[triali, selcue] = Q_est[triali, selcue] + \
            alpha[triali, selcue] * rpe
        simDict['RPE_H'].append(rpe)
        simDict['alpha_H'].append(alpha[triali, selcue])
        for qi, q_est in enumerate(Q_est[triali, :]):
            simDict[f'Qest_{qi}_H'].append(q_est)

        # Response time
        simDict['rt_H'].append(sim_rt(rpe))

    return af.save_data(simDict, data, var_list)

def simMeta_1c(parameters, data, asm='soft'):
    """
    Michelle the meta learner
    The meta learning model of Cohen and his team.
    https://doi.org/10.1016/j.neuron.2019.06.001
    https://doi.org/10.1016/j.cub.2021.12.006

    Parameters
    ----------
    parameters : tuple, list, array
        First parameter is the learning rate of the model (0 <= alpha <= 1).
        Second parameter is the forgetting rate of the model (zeta).
        Third parameter is the constant of the action selection method
            (0 < beta).
        Fourth parameter is the bias term of the biased SoftMax
            (-1 (cue_0 bias) <= bias <= 1 (cue_1 bias)). It indicates a prefere
            for one cue over the other. The default is 0 (no bias).
    data : pandas.DataFrame
        Dataframe containing the structure of the experiment.
    asm : string, optional
        The action selection method, policy. The default is SoftMax.

    Returns
    -------
    simData : pandas.DataFrame
        Contains columns of 'data' and simulated behaviour of Michelle:
            0. 'selCue_M' - selected cue by the meta learning model
            1. 'prob_M' - probability to select cue 0
            2. 'rt_M' - response time
            3. 'reward_M' - reward based on cue selected by meta learner
            4. 'RPE_M' - reward prediction error
            5. 'Qest_0_M' - estimated value of cue 0
            6. 'Qest_1_M' - estimated value of cue 1
            7. 'eun_0_M' - estimate of expected uncertainty of cue 0
            8. 'eun_1_M' - estimate of expected uncertainty of cue 1
            9. 'uun_M' - unexpected uncertainty
    """
    # Variables
    # ---------
    # Dataframe
    var_list = [
        'selCue_M', 'prob_M', 'rt_M', 'reward_M',
        'RPE_M', 'Qest_0_M', 'Qest_1_M',
        'eun_0_M', 'eun_1_M', 'uun_M'
        ]
    simDict = {vari:[] for vari in var_list}

    # Number of trials
    n_trials = data.shape[0]
    # Number of cues
    N_CUES = 2
    # Model
    ## Estimated value of cue
    Q_est = np.full((n_trials, N_CUES), 1/N_CUES)
    ## Outcomes
    outcomes = np.full((2, ), np.nan)
    ## Estimate of expected uncertainty calculated from the history of URPEs
    eun = np.full((n_trials, ), 1/N_CUES)
    ## Unexpected uncertainty
    uun = 0

    # Policy
    ## Selected cue
    selcue = np.nan

    # Parameters
    alpha, zeta = parameters[:2]
    beta, bias = 0, 0
    if len(parameters) > 2:
        beta = parameters[2]
    if len(parameters) > 3:
        bias = parameters[3]

    # Trial loop
    for triali, trial in data.iterrows():
        # Policy
        # ------
        ## Random for first trial
        if triali == 0:
            selcue = np.random.randint(N_CUES)
            probcue = 0.5
        else:
            selcue, probcue = af.policy(asm, Q_est[triali - 1, :],
                                        beta, bias)
        simDict['selCue_M'].append(selcue)
        simDict['prob_M'].append(probcue)

        # Reward
        # ------
        ## Based on validity
        ##AM: If cue==target reward = 1, if cue!=target reward = 0
        outcomes[int(trial.relCueCol)] = trial.relCue == trial.targetLoc
        outcomes[int(1 - trial.relCueCol)] = trial.irrelCue == trial.targetLoc
        reward = int(outcomes[selcue])
        simDict['reward_M'].append(reward)

        # Update rule
        # -----------
        if triali == 0:
            # Reward prediction error
            rpe = reward - Q_est[triali, selcue]
        else:
            # Reward prediction error
            rpe = reward - Q_est[triali - 1, selcue]
            # Cue estimates
            ## Repeat values of previous trial
            Q_est[triali, :] = Q_est[triali - 1, :]
            eun[trial] = eun[trial - 1]
        ## Update expected and unexpected uncertainty
        uun = abs(rpe) - eun[triali]
        eun[triali] *= uun
        ## Update cue estimate of selected stimulus in current trial
        Q_est[triali, selcue] = Q_est[triali, selcue] + \
            alpha * rpe * (1 - eun[triali])
        ## Forget not selected stimulus
        Q_est[triali, abs(1 - selcue)] = zeta * Q_est[triali, abs(1 - selcue)]
        simDict['RPE_M'].append(rpe)
        for qi, q_est in enumerate(Q_est[triali, :]):
            simDict[f'Qest_{qi}_M'].append(q_est)

        # Response time
        simDict['rt_G'].append(sim_rt(rpe))

    return af.save_data(simDict, data, var_list)

#%% ~~ Models (2 cues) ~~ %%#
#---------------------------#

# ~~ Rescorla - Wagner ~~ #
def simRW_2c(parameters, data, asm='soft'):
    """
    Daphne the delta learner
    The Rescorla - Wagner learning model with SoftMax. The policy calculations
    are part of the function (and not a separate function). The cue
    estimates of both cues are updated.

    Parameters
    ----------
    parameters : tuple, list, array
        First parameter is the learning rate of the model (0 <= alpha <= 1).
        Second parameter is the constant of the action selection method
            (0 < beta).
    data : pandas.DataFrame
        Dataframe containing the structure of the experiment.
    asm : string, optional
        The action selection method, policy. The default is SoftMax.

    Returns
    -------
    simData : pandas.DataFrame
        Contains columns of 'data' and simulated behaviour of Daphne:
            0. 'selCue_RW2' - selected cue by the Rescorla-Wagner model
            1. 'prob_RW2' - probability to select cue 0
            2. 'rt_RW2' - response time
            3. 'reward_RW2' - reward based on cue selected by Rescorla-Wagner
            4. 'RPE_RW2' - reward prediction error
            5. 'Qest_0_RW2' - estimated value of cue 0
            6. 'Qest_1_RW2' - estimated value of cue 1
    """

    # Variables
    # ---------
    # Dataframe
    var_list = [
        'selCue_RW2', 'prob_RW2', 'rt_RW2', 'reward_RW2',
        'RPE_RW2', 'Qest_0_RW2', 'Qest_1_RW2'
        ]
    simDict = {vari:[] for vari in var_list}

    ## Number of trials
    n_trials = data.shape[0]
    ## Number of cues
    N_CUES = 2
    # Model
    ## Estimated value of cue
    Q_est = np.full((n_trials, N_CUES), 1/N_CUES)
    ## Outcomes
    outcomes = np.full((2, ), np.nan)

    # Policy
    ## Selected cue
    selcue = np.nan

    # Parameters
    alpha = parameters[0]
    beta, bias = 0, 0
    if len(parameters) > 1:
        beta = parameters[1]
    if len(parameters) > 2:
        bias = parameters[2]

    # Trial loop
    for triali, trial in data.iterrows():
        # Policy
        # ------
        ## Random for first trial
        if triali == 0:
            selcue = np.random.randint(N_CUES)
            probcue = 0.5
        else:
            selcue, probcue = af.policy(asm, Q_est[triali - 1, :],
                                        beta, bias)
        simDict['selCue_RW2'].append(selcue)
        simDict['prob_RW2'].append(probcue)

        # Reward
        # ------
        ## Based on validity
        ##AM: If cue==target reward = 1, if cue!=target reward = 0
        outcomes[int(trial.relCueCol)] = trial.relCue == trial.targetLoc
        outcomes[int(1 - trial.relCueCol)] = trial.irrelCue == trial.targetLoc
        reward = int(outcomes[selcue])
        simDict['reward_RW2'].append(reward)

        # Update rule (RW)
        # ----------------
        if triali == 0:
            for cuei in range(N_CUES):
                # Reward prediction error
                rpe = reward - Q_est[triali, cuei]
                # Cue estimates
                Q_est[triali, cuei] = Q_est[triali, cuei] + alpha * rpe
        else:
            # Repeat cue estimates of previous trial
            Q_est[triali, :] = Q_est[triali - 1, :]
            for cuei in range(N_CUES):
                # Reward prediction error
                rpe = reward - Q_est[triali - 1, cuei]
                # Cue estimates
                Q_est[triali, cuei] = Q_est[triali, cuei] + alpha * rpe
        simDict['RPE_RW2'].append(reward - Q_est[triali - 1, selcue])
        for qi, q_est in enumerate(Q_est[triali, :]):
            simDict[f'Qest_{qi}_RW2'].append(q_est)

        # Response time
        simDict['rt_RW2'].append(sim_rt(rpe))

    return af.save_data(simDict, data, var_list)

# ~~ RW-PH Hybrid ~~ #
def simHybrid_2c(parameters, data, salpha=0.01, asm='soft'):
    """
    Hugo the hybrid learner
    The RW-PH hybrid learning model with SoftMax. The policy calculations
    are part of the function (and not a separate function). The cue
    estimates of both cues are updated.

    Parameters
    ----------
    parameters : tuple, list, array
        First parameter is eta (0 <= eta <= 1).
        Second parameter is the constant of the action selection method
            (0 < beta).
    data : pandas.DataFrame
        Dataframe containing the structure of the experiment.
    salpha : float, optional
        Start of alpha. The default is 0.01.
    asm : string, optional
        The action selection method, policy. The default is SoftMax.

    Returns
    -------
    simData : pandas.DataFrame
        Contains columns of 'data' and simulated behaviour of Hugo:
            0. 'selCue_H2' - selected cue by the hybrid model
            1. 'prob_H2' - probability to select cue 0
            2. 'rt_H2' - response time
            3. 'reward_H2' - reward based on cue selected by hybrid model
            4. 'alpha_H2' - learning rate
            5. 'RPE_H2' - reward prediction error
            6. 'Qest_0_H2' - estimated value of cue 0
            7. 'Qest_1_H2' - estimated value of cue 1
    """

    # Variables
    #----------
    # Dataframe
    var_list = [
        'selCue_H2', 'prob_H2', 'rt_H2', 'reward_H2', 'alpha_H2',
        'RPE_H2', 'Qest_0_H2', 'Qest_1_H2'
        ]
    simDict = {vari:[] for vari in var_list}

    ## Number of trials
    n_trials = data.shape[0]
    ## Number of cues
    N_CUES = 2

    # Model
    ## Alpha
    alpha = np.full((n_trials, N_CUES), salpha)
    ## Estimated value of cue
    Q_est = np.full((n_trials, N_CUES), 1/N_CUES)
    ## Outcomes
    outcomes = np.full((2, ), np.nan)

    # Policy
    ## Selected cue
    selcue = np.nan

    # Parameters
    eta = parameters[0]
    beta, bias = 0, 0
    if len(parameters) > 1:
        beta = parameters[1]
    if len(parameters) > 2:
        bias = parameters[2]

    # Trial loop
    for triali, trial in data.iterrows():
        # Policy
        # ------
        ## Random for first trial
        if triali == 0:
            selcue = np.random.randint(N_CUES)
            probcue = 0.5
        else:
            selcue, probcue = af.policy(asm, Q_est[triali - 1, :],
                                        beta, bias)
        simDict['selCue_H2'].append(selcue)
        simDict['prob_H2'].append(probcue)

        # Reward
        # ------
        ## Based on validity
        ##AM: If cue==target reward = 1, if cue!=target reward = 0
        outcomes[int(trial.relCueCol)] = trial.relCue == trial.targetLoc
        outcomes[int(1 - trial.relCueCol)] = trial.irrelCue == trial.targetLoc
        reward = int(outcomes[selcue])
        simDict['reward_H2'].append(reward)

        # Hybrid
        # ------
        if triali == 0:
            for cuei in range(N_CUES):
                # Reward prediction error
                rpe = reward - Q_est[triali, cuei]
                # Alpha (PH)
                alpha[triali, cuei] = eta * np.abs(rpe) + \
                    (1 - eta) * alpha[triali, cuei]
                # Cue estimates
                Q_est[triali, cuei] = Q_est[triali, cuei] + \
                    alpha[triali, cuei] * rpe
        else:
            ## Repeat values of previous trial
            Q_est[triali, :] = Q_est[triali - 1, :]
            for cuei in range(N_CUES):
                # Reward prediction error
                rpe = reward - Q_est[triali - 1, cuei]
                # Alpha (PH)
                alpha[triali, cuei] = eta * np.abs(rpe) + \
                    (1 - eta) * alpha[triali, cuei]
                # Cue estimates
                Q_est[triali, cuei] = Q_est[triali, cuei] + \
                    alpha[triali, cuei] * rpe
        simDict['RPE_H2'].append(reward - Q_est[triali - 1, selcue])
        simDict['alpha_H2'].append(alpha[triali, selcue])
        for qi, q_est in enumerate(Q_est[triali, :]):
            simDict[f'Qest_{qi}_H2'].append(q_est)

        # Response time
        simDict['rt_H2'].append(sim_rt(rpe))

    return af.save_data(simDict, data, var_list)

#%% ~~ Controle models ~~ %%#
#---------------------------#

# ~~ Wilhelm ~~ #
def simWSLS(data):
    """
    Wilhelm the Win-stay-lose-shift model
    Selects the cue that was predictive during its previous presentation.

    Parameters
    ----------
    data : pd.DataFrame
        Dataframe containing the structure of the experiment.

    Returns
    -------
    simData : pd.DataFrame
        Contains columns of 'data' and simulated behaviour of Wilhelm:
            0. 'selCue_W' - selected cue by the WSLS model
            1. 'prob_W' - probability to select cue 0
            2. 'reward_W' - reward based on cue selected by WSLS model
    """

    # Variables
    # ---------
    # Dataframe
    var_list = ['selCue_W', 'prob_W', 'reward_W']
    simDict = {vari:[] for vari in var_list}

    # Model
    ## Number of trials
    n_trials = data.shape[0]
    ## Number of cues
    N_CUES = 2
    ## Outcomes
    outcomes = np.full((2, ), np.nan)
    ## Reward
    reward = np.nan
    ## Selected cues
    selcues = np.full((n_trials, ), np.nan)

    # Trial loop
    for triali, trial in data.iterrows():
        # Policy
        # ------
        ## Random for first trial
        if triali == 0:
            selcues[triali] = np.random.randint(N_CUES)
        else:
            if reward == 1:
                selcues[triali] = selcues[triali - 1]
            else:
                selcues[triali] = 1 - selcues[triali - 1]
        # Reward calculations
        ## Based on validity
        ##AM: If cue==target reward = 1, if cue!=target reward = 0
        outcomes[int(trial.relCueCol)] = trial.relCue == trial.targetLoc
        outcomes[int(1 - trial.relCueCol)] = trial.irrelCue == trial.targetLoc
        reward = int(outcomes[int(selcues[triali])])

        simDict['selCue_W'].append(selcues[triali])
        ##SG: If there is some uncertainty in which cue Wilhelm will pick,
            # there is something wrong in the model.
        simDict['prob_W'].append(1)
        simDict['reward_W'].append(reward)

    return af.save_data(simDict, data, var_list)

# ~~ Renee ~~ #
def simRandom(data):
    """
    Renee the random model
    Random cue selection.

    Parameters
    ----------
    data : pd.DataFrame
        Dataframe containing the structure of the experiment.

    Returns
    -------
    simData : pd.DataFrame
        Contains columns of 'data' and simulated behaviour of Renee:
            0. 'selCue_R' - selected cue by the random model
            1. 'prob_R' - probability to select cue 0
            2. 'reward_R' - reward based on cue selected by random model
    """

    # Variables
    # ---------
    # Dataframe
    var_list = ['selCue_R', 'prob_R', 'reward_R']
    simDict = {vari:[] for vari in var_list}

    # Number of cues
    N_CUES = 2
    ## Outcomes
    outcomes = np.full((2, ), np.nan)

    # Trial loop
    for triali, trial in data.iterrows():
        selcue = np.random.randint(N_CUES)
        simDict['selCue_R'].append(selcue)
        simDict['prob_R'].append(0.5)
        outcomes[int(trial.relCueCol)] = trial.relCue == trial.targetLoc
        outcomes[int(1 - trial.relCueCol)] = trial.irrelCue == trial.targetLoc
        reward = int(outcomes[selcue])
        simDict['reward_R'].append(reward)

    return af.save_data(simDict, data, var_list)

#%% ~~ Others ~~ %%#
####################

def sim_negLL(thetas, data, model):
    """
    Negative log-likelihood of choice based on RW predictions of simulations

    Parameters
    ----------
    thetas : list, array, tuple
        Parameter values.
    data : pd.DataFrame
        Prepared data of simulation, containing structure of experiment.
        Doesn't need to go through rescon first.
    model : string
        Name of the used model:
            RW - Rescorla-Wagner
            H - RW-PH hybrid
            W - Win-stay-lose-shift
            R - Random

    Returns
    -------
    Negative log-likelihood of selected stimuli.
    """

    model = model.upper()
    # Rename to avoid duplicates
    data = data.rename(columns={f'selCue_{model}': 'selCue',
                                f'prob_{model}': 'prob'})
    # Log-likelihood of choice
    loglik_of_choice = []
    # Simulate data
    modelDict = sim_models()
    if not model in ['W', 'R']:
        simData = modelDict[model](thetas, data)
    else:
        simData = modelDict[model](data)

    for _, trial in simData.iterrows():
        # Model estimated value of model's selected stimulus
        ##SG: Likelihood of left stimulus if model picked left.
            # Otherwise 1-left for likelihood of right stimulus.
        picked_prob = abs(trial[f'selCue_{model}'] -
                          trial[f'prob_{model}'])
        loglik_of_choice.append(np.log(picked_prob))

    # Return negative log likelihood
    # return -1 * np.nansum(loglik_of_choice)
    return np.nansum(loglik_of_choice)

def sim_negSpearCor(thetas, data, model, asm='soft'):
    """
    Negative Spearman correlation of learning model of reaction time and
    estimated value of selected cue

    Parameters
    ----------
    thetas : list, array, tuple
        Parameter values.
    data : pd.DataFrame
        Prepared data of simulation, containing structure of experiment.
        Doesn't need to go through rescon first.
    model : string
        Name of the used model:
            RW - Rescorla-Wagner
            H - RW-PH hybrid
            M - Meta learner
    asm : string, optional
        The action selection method, policy. The default is SoftMax.

    Returns
    -------
    negative spearman r
    """

    model = model.upper()
    # Wilhelm and Renee do not simulate RTs.
    assert not model in ['W', 'R'], 'Model has no simulated RT'

    # Rename to avoid duplicates
    data = data.rename(columns={f'selCue_{model}': 'selCue',
                                f'prob_{model}': 'prob',
                                f'rt_{model}': 'rt',
                                f'RPE_{model}': 'RPE'})
    # Simulate data
    modelDict = sim_models()
    simData = modelDict[model](parameters=thetas, data=data, asm=asm)

    # Correlation between RT and RPE
    return - stats.spearmanr(simData[f'rt_{model}'].to_numpy(),
                             # simData[f'RPE_{model}'].to_numpy(),
                             abs(simData[f'RPE_{model}']).to_numpy(),
                             nan_policy = 'omit')[0]

def sim_accuracy(thetas, data, model, asm='soft'):
    """
    Model accuracy

    Parameters
    ----------
    thetas : list, array, tuple
        Parameter values.
    data : pd.DataFrame
        Prepared data of simulation, containing structure of experiment.
        Doesn't need to go through rescon first.
    model : string
        Name of the used model:
            RW - Rescorla-Wagner
            H - RW-PH hybrid
            M - Meta learner
            W - Win-stay-lose-shift
            R - Random
    asm : string, optional
        The action selection method, policy. The default is SoftMax.

    Returns
    -------
    Accuracy
    """

    # Rename to avoid duplicates
    data = data.rename(columns={f'selCue_{model}': 'selCue',
                                f'prob_{model}': 'prob',
                                f'rt_{model}': 'rt',
                                f'RPE_{model}': 'RPE'})

    model = model.upper()
    # Simulate data
    modelDict = sim_models()
    if model in ['W', 'R']:
        simData = modelDict[model](data)
    else:
        simData = modelDict[model](parameters=thetas, data=data, asm=asm)

    return sum(simData['relCue'] == simData[f'selCue_{model}']) / len(data)

def var_bin_switch(
        dataList, datadir, varList, validity='validity', bin_size=15, uun='All'
        ):
    """
    Bin variability effect of simulated data before and after switch

    Parameters
    ----------
    dataList : list
        List containing the data filenames.
    datadir : Path
        Location of the data.
    varList : tuple, list, array
        List with the names of the relevant variables.
    bin_size : int, optional
        Number of trials grouped together. The default is 15.

    Returns
    -------
    post15 : dictionary
        First 15 trials after switch.
    middle15 : dictionary
        Trials 15 to 30.
    leftover : dictionary
        All trials except for first 30.
    last15 : dictionary
        Last 15 trials, less if there are less then 45 trials between switches.
    pre15 : dictionary
        The 15 trials before switch.
    """

    # Trial bins
    post15 = {vari:[] for vari in varList}
    middle15 = {vari:[] for vari in varList}
    leftover = {vari:[] for vari in varList}
    last15 = {vari:[] for vari in varList}
    pre15 = {vari:[] for vari in varList}

    for simi, filei in enumerate(dataList):
        sim_data = pd.read_csv(datadir / filei, index_col='Unnamed: 0')
        if uun.upper() == 'HIGH':
            sim_data = sim_data[sim_data['gammaBlock'] <= 0.8]
        elif uun.upper() == 'LOW':
            sim_data = sim_data[sim_data['gammaBlock'] > 0.8]
        sim_data.reset_index(drop=True, inplace=True)

        # Switch points
        lag_relCueCol = sim_data.relCueCol.eq(sim_data.relCueCol.shift())
        switch_points = np.where(lag_relCueCol == False)[0]
        switch_points = np.append(switch_points, len(sim_data))

        for starti, endi in af.pairwise(switch_points):
            nover = endi - starti - (2 * bin_size)
            for vari in varList:
                # First 15 trials after switch
                post15_data = sim_data.loc[starti:
                                           (starti + bin_size - 1)][[validity, vari]]
                post15[vari].append(
                    np.nanmean(post15_data[post15_data[validity] == 0][vari]) -\
                    np.nanmean(post15_data[post15_data[validity] == 1][vari])
                    )
                # Trials 15 to 30
                middle15_data = sim_data.loc[(starti + bin_size):
                                             (starti + 2 * bin_size - 1)][[validity, vari]]
                middle15[vari].append(
                    np.nanmean(middle15_data[middle15_data[validity] == 0][vari]) -\
                    np.nanmean(middle15_data[middle15_data[validity] == 1][vari])
                    )
                # All trials except for first 30
                left_data = sim_data.loc[(starti + 2 * bin_size):
                                         (endi - 1)][[validity, vari]]
                leftover[vari].append(
                    np.nanmean(left_data[left_data[validity] == 0][vari]) -\
                    np.nanmean(left_data[left_data[validity] == 1][vari])
                    )
                # The 15 trials before switch
                pre15_data = sim_data.loc[(endi - bin_size):
                                          (endi - 1)][[validity, vari]]
                pre15[vari].append(
                    np.nanmean(pre15_data[pre15_data[validity] == 0][vari]) -\
                    np.nanmean(pre15_data[pre15_data[validity] == 1][vari])
                    )
                # Last 15 trials or less if there were less than 45 trials
                # between switches
                if nover >= bin_size:
                    last15_data = sim_data.loc[(endi - bin_size):
                                               (endi - 1)][[validity, vari]]
                else:
                    last15_data = sim_data.loc[(endi - nover):
                                               (endi - 1)][[validity, vari]]
                last15[vari].append(
                    np.nanmean(last15_data[last15_data[validity] == 0][vari]) -\
                    np.nanmean(last15_data[last15_data[validity] == 1][vari])
                    )

    return post15, middle15, leftover, last15, pre15

#------------------------------------------------------------------------- End


import os
import numpy as np
import log
import argparse
import pandas as pd
import math
import random
import warnings
warnings.filterwarnings('ignore')
#日志打印到屏幕并保存到文件
if not os.path.exists('result'):
    os.mkdir('./result')
logger = log.Logger(filename='result/info.log').logger

class Node:
    """创建一个树的节点"""
    def __init__(self,
                 data_index,
                 split_feature=None,
                 split_value=None,
                 is_leaf=False,
                 loss=None,
                 current_depth=None):
        """
        :param data_index:该节点的数据在全部数据集中的索引
        :param split_feature: 最佳分割特征
        :param split_value:最佳分割特征值
        :param is_leaf:是否为叶子节点
        :param loss: 分类损失
        :param current_depth: 当前节点所在树的深度
        """
        self.loss = loss
        self.split_feature = split_feature
        self.split_value = split_value
        self.data_index = data_index
        self.is_leaf = is_leaf  #是不是叶子节点
        self.predict_value = None
        self.left_child = None  #如果当前节点可继续划分，则分为左子树和右子树
        self.right_child = None
        self.current_depth = current_depth

    def update_predict_value(self, targets, sample_weight=None):
        self.predict_value = self.loss.update_leaf_values(targets, sample_weight=sample_weight)
        logger.info('>>>>>>>>>>>叶子节点预测值：%.3f'%self.predict_value)

    def get_predict_value(self, instance):
        """
        预测结果，采用递归的方法
        :param instance: 一个样本
        :return:
        """
        if self.is_leaf:
            #如果是叶子节点，直接获得叶子节点值
            return self.predict_value
        if instance[self.split_feature] < self.split_value:
            return self.left_child.get_predict_value(instance)
        else:
            return self.right_child.get_predict_value(instance)

class Tree:
    """创建一棵树，从根节点开始创建"""
    def __init__(self, X, y, max_depth=3, min_samples_split=2, features_name=None, loss=None, sample_weight=None):
        """
        初始化树的参数
        :param data:
        :param max_depth:树的最大深度
        :param min_samples_split:最小划分数据量
        :param sampel_weight: 样本权重
        :param features:特征值
        :param loss:
        :return: 树结构和节点
        """
        self.loss = loss
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.sample_weight=np.ones(len(X)) if sample_weight is None else sample_weight
        self.features_name = features_name
        self.remain_index = np.array([i for i in range(len(X))])  #当前节点的样本在原始数据中的下标索引
        self.leaf_nodes = []
        self.root_node = self.build_tree(X, y, self.remain_index, depth=0, sample_weight=self.sample_weight)  #根节点

    def build_tree(self, X, y, remain_index, depth=0, sample_weight=None):
        """
        构建一棵树
        此处有三个树继续生长的条件：
        1: 深度没有到达最大, 树的深度假如是3， 意思是需要生长成3层, 那么这里的depth只能是0, 1，所以判断条件是 depth < self.max_depth - 1
        2: 点样本数 >= min_samples_split
        3: 此节点上的样本的 target_name 值不一样（如果值 一样说明已经划分得很好了，不需要再分）
        :param sample_weight:
        :param X: 原始样本特征值
        :param y: 原始样本标签值
        :param remain_index:当前数据在原始数据中的索引
        :param depth: 树的当前深度
        :return: 树节点
        """
        now_X, now_y = X[remain_index], y[remain_index]
        n_samples, n_features=now_X.shape[0], now_X.shape[1]
        if depth < self.max_depth and n_samples >= self.min_samples_split and len(np.unique(now_y)) > 1:
            se = None
            split_feature = None
            split_value = None
            left_index_of_now_data = None
            right_index_of_now_data = None
            left_weight_index_of_now_data=None
            right_weight_index_of_now_data=None
            logger.info('----------树的深度：%d----------' % depth)

            """遍历每一个特征，分别计算最优的切割点"""
            for i in range(n_features):
                #logger.info('----------划分特征：%s----------'%(self.features_name[i]))
                #对特征进行预排序,获得是特征值的下标索引
                presort=np.argsort(now_X[:,i])
                '''遍历每个特征值的每个样本值'''
                for j in range(n_samples):
                    # 尝试划分
                    left_index =presort[:j]
                    right_index =presort[j:]
                    left_se = calculate_se(y[left_index], sample_weight=sample_weight[left_index])
                    right_se = calculate_se(y[right_index], sample_weight=sample_weight[right_index])
                    sum_se = left_se + right_se

                    if se is None or sum_se < se:
                        split_feature = i
                        split_value=X[remain_index[right_index[0]],i]
                        se = sum_se
                        left_index_of_now_data = remain_index[left_index]
                        right_index_of_now_data = remain_index[right_index]
                        left_weight_index_of_now_data=left_index
                        right_weight_index_of_now_data=right_index
            logger.info('----------最佳划分特征：%s----------'%split_feature)
            logger.info('----------最佳划分值：%.3f----------'%split_value)

            #找到当前最优划分特征和划分值，对当前这个节点进行处理，产生子节点
            node = Node(remain_index, split_feature, split_value, current_depth=depth)

            logger.info('----------构建左子树----------')
            node.left_child = self.build_tree(X, y, left_index_of_now_data, depth=depth + 1, sample_weight=sample_weight[left_weight_index_of_now_data])
            logger.info('----------构建右子树----------')
            node.right_child = self.build_tree(X, y, right_index_of_now_data, depth=depth + 1, sample_weight=sample_weight[right_weight_index_of_now_data])
            return node
        else:
            #不满足分裂条件，停止分裂，该节点就是叶子节点
            logger.info('----------树的深度：%d----------' % depth)
            node = Node(remain_index, is_leaf=True, loss=self.loss, current_depth=depth)
            node.update_predict_value(y[remain_index], sample_weight=sample_weight)
            self.leaf_nodes.append(node)
            return node

class BinomialDeviance:
    """二分类"""
    def initialize_f_0(self, y):
        """初始化 F_0 二分类的损失函数是对数损失，初始值是正样本的个数与负样本个数的比值取对数"""
        pos = y.sum()   #正样本
        neg = len(y) - pos   #负样本

        f_0_val = math.log(pos / neg)
        return np.ones(len(y))*f_0_val

    def gradients_(self,y_pred, y_true):
        """一阶导数"""
        return 1 / (1 + np.exp(-y_pred)) - y_true

    def hessians_(self, y_pred, y_true=None):
        """二阶导数"""
        return np.exp(-y_pred) / ((1 + np.exp(-y_pred)) ** 2)

    def calculate_residual(self, y, current_pred_value, iter):
        """计算负梯度"""
        residual=-self.gradients_(current_pred_value, y)
        return residual

    def update_f_m(self, X, current_pred_value, trees, iter, learning_rate):
        """计算 当前时刻的预测值F_m ，"""
        for leaf_node in trees[iter].leaf_nodes:
            current_pred_value[leaf_node.data_index]+=learning_rate*leaf_node.predict_value
        return current_pred_value

    def update_leaf_values(self, targets, sample_weight=None):
        """更新叶子节点的预测值"""
        if len(targets)!=0:
            return (targets*sample_weight).sum()/sample_weight.sum()
        else:
            return 0.0

    def get_train_loss(self, y, f, iter,):
        """计算训练损失"""
        loss = -2.0 * ((y * f) - f.apply(lambda x: math.exp(1+x))).mean()
        logger.info(('第%d棵树: log-likelihood:%.4f' % (iter, loss)))

class BaseGradientBoosting():
    def __init__(self, loss, learning_rate=0.1, n_estimators=100, max_depth=3, min_samples_split=2):
        super().__init__()
        self.loss = loss  #损失
        self.learning_rate = learning_rate  #学习率
        self.n_estimators = n_estimators  #树的棵数
        self.max_depth = max_depth   #单个树最大深度
        self.min_samples_split = min_samples_split #单节点最小划分样本量
        self.features_name=None
        self.trees = {}
        self.f_0 = {}

    def fit(self, X, y, sample_weight=None):
        """
        训练树模型
        :param X: 特征值
        :param y: 样本label真实值
        :param sample_weight: 样本权重，默认为None
        :return:
        """
        #TODO 删除id和label，得到特征名称，这里的特征名称是加密的，直接用代号就可以了
        self.features_name = ['fea_'+str(i) for i in range(X.shape[1]) ]

        # 初始化 f_0(x),初始化的计算方式取决于损失函数，对于（分类）交叉熵损失--来说，初始化 f_0(x) 就是 log(正/负)；
        self.f_0 = self.loss.initialize_f_0(y)

        # 对 m = 1, 2, ..., M
        current_pred_value=self.f_0  #当前树的预测值
        for iter in range(1, self.n_estimators+1): #遍历所有的树节点
            # 计算负梯度--对于平方误差来说就是残差
            logger.info('-----------------------------构建第%d颗树-----------------------------' % iter)
            residual= self.loss.calculate_residual(y, current_pred_value, iter=iter) #计算负梯度
            self.trees[iter] = Tree(X, residual, self.max_depth, self.min_samples_split, self.features_name, self.loss)
            current_pred_value=self.loss.update_f_m(X, current_pred_value, self.trees, iter, self.learning_rate)

class GradientBoostingBinaryClassifier(BaseGradientBoosting):
    def __init__(self,learning_rate=0.1,n_estimators=100,max_depth=3,min_samples_split=2):
        super().__init__(BinomialDeviance(),learning_rate,n_estimators,max_depth,min_samples_split)

    def predict(self, X, return_proba=False):
        f_0=self.f_0
        score= np.zeros(X.shape[0])
        for iter in range(1, self.n_estimators + 1):
             for k in range(X.shape[0]):
                 score[k]+=self.trees[iter].root_node.get_predict_value(X[k])

        # 相比于回归任务，分类任务需把要最后累加的结果Fm(x)转成概率
        predict_proba=1/(1+np.exp(-score))
        if return_proba:
            return predict_proba
        score[predict_proba>=0.5]=1
        score[predict_proba<0.5]=0
        predict_label=score
        return predict_label

def calculate_se(label,sample_weight=None):
    """计算平方误差 mse"""
    mean = (label*sample_weight).sum()/sample_weight.sum()
    se = 0
    for y in label:
        se += (y - mean) * (y - mean)
    return se

def load_data():
    #训练数据和测试数据的demo
    from sklearn.datasets import load_breast_cancer
    x,y=load_breast_cancer(return_X_y=True)
    # data=[[1, 5, 20, 0],
    #       [2, 21, 70, 1],
    #       [3, 7, 30, 1],
    #       [4, 30, 60, 0]]
    # x=np.array(data)[:,[1,2]]
    # y=np.array(data)[:,-1]

    return x,y

def run(args):
    from sklearn.model_selection import train_test_split
    from sklearn.metrics import accuracy_score
    # 获取训练和测试数据
    X,y = load_data()
    x_train,x_test,y_train,y_test=train_test_split(X,y,random_state=20191126,test_size=0.1)

    model = GradientBoostingBinaryClassifier(learning_rate=args.lr,
                                             n_estimators=args.trees,
                                             max_depth=args.depth)
    # 训练模型
    model.fit(x_train,y_train)
    # 模型预测
    y_pred=model.predict(x_test,return_proba=False)
    print(y_pred)
    acc=accuracy_score(y_test,y_pred)
    print(acc)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--lr', default=0.1, type=float, help='learning rate')
    parser.add_argument('--trees', default=5, type=int, help='the number of decision trees')
    parser.add_argument('--depth', default=3, type=int, help='the max depth of decision trees')
    # 非叶节点的最小数据数目，如果一个节点只有一个数据，那么该节点就是一个叶子节点，停止往下划分
    parser.add_argument('--count', default=2, type=int, help='the min data count of a node')
    parser.add_argument('--log', default=False, type=bool, help='whether to print the log on the console')
    parser.add_argument('--plot', default=True, type=bool, help='whether to plot the decision trees')
    args = parser.parse_args()
    run(args)