import os
import numpy as np
import pandas as pd
from scipy import stats

import opensim as osim

NIND = 200 # Population size
GEN = 15 # Generation number
TIMEOUT = 900 # Timeout for each simulation

LOG_LEVEL = {
    'DEBUG': 1,
    'INFO': 2,
    'WARN': 3,
    'ERROR': 4
}

NO_MUSCLES = {
    'thorax': 7,
    'coxa': 5,
    'femur': 2,
}


def update_coordinates_in_opensim(xml_path, new_coordinates_file, output_path, final_time=0.4):
    # Load the AnalyzeTool from the XML
    tool = osim.AnalyzeTool(xml_path)
    # Update the coordinates file and save the updated settings
    tool.setCoordinatesFileName(new_coordinates_file)
    # Update final time
    tool.setFinalTime(final_time)
    # CHange analysisset objects staticoptimization end_time
    
    # tool.AnalysisSet().get(0).end_time = final_time 

    tool.printToXML(output_path)
    print(f"Coordinates file updated in {output_path}")

def calc_rmse(data1, data2):
    """Calculate the root mean square error between two data sets."""
    return np.sqrt(np.mean((data1 - data2) ** 2))

def calc_pearson(data1, data2):
    """Calculate the Pearson correlation coefficient between two data sets."""
    return stats.pearsonr(data1, data2).statistic

def vec_to_list(vec):
    """Converts an OpenSim Vec3/6 to a Python list."""
    return [vec.get(i) for i in range(vec.size())]

def convert_time(pandas_data):
    """ Converts pandas data to a valid format. """
    multiply_w = 1000000
    pandas_data['time'] = pd.to_datetime((pandas_data['time']*multiply_w).apply(int).astype('str').str.zfill(6), format='%f')
    pandas_data = pandas_data.set_index('time')
    return pandas_data


def read_motion_file(filename):
    """Reads OpenSim .sto files and returns the data as a DataFrame.

    Parameters
    ----------
    filename: str
        Absolute path to the .sto file.

    Returns
    -------
    header: list
        The header of the .sto file.
    data: pd.DataFrame
        DataFrame containing all motion data with proper column names.
    """

    if not os.path.exists(filename):
        raise FileNotFoundError(f"File '{filename}' does not exist.")

    with open(filename, 'r') as file_id:
        # Read header
        header = []
        nc, nr = 0, 0
        while True:
            next_line = file_id.readline().strip()
            header.append(next_line)
            if 'endheader' in next_line:
                break
            if 'nColumns' in next_line:
                nc = int(next_line.split('=')[-1].strip())
            if 'nRows' in next_line:
                nr = int(next_line.split('=')[-1].strip())

        # Read column labels
        labels = file_id.readline().strip().split()

        # Read data
        data = pd.read_csv(file_id, sep='\s+', names=labels)

    return header, data


def readMotionFile(filename,cols):
    """ Reads OpenSim .sto files.
    Parameters
    ----------
    filename: absolute path to the .sto file
    Returns
    -------
    header: the header of the .sto
    labels: the labels of the columns
    data: an array of the data
    """

    if not os.path.exists(filename):
        print('file do not exists')

    file_id = open(filename, 'r')

    # read header
    next_line = file_id.readline()
    header = [next_line]
    nc = 0
    nr = 0
    while not 'endheader' in next_line:
        if 'datacolumns' in next_line:
            nc = int(next_line[next_line.index(' ') + 1:len(next_line)])
        elif 'datarows' in next_line:
            nr = int(next_line[next_line.index(' ') + 1:len(next_line)])
        elif 'nColumns' in next_line:
            nc = int(next_line[next_line.index('=') + 1:len(next_line)])
        elif 'nRows' in next_line:
            nr = int(next_line[next_line.index('=') + 1:len(next_line)])

        next_line = file_id.readline()
        header.append(next_line)

    # process column labels
    next_line = file_id.readline()
    if next_line.isspace() == True:
        next_line = file_id.readline()

    labels = next_line.split()
    # get data
    data = []
    for i in range(1, nr + 1):
        # d = [float(x) for x in file_id.readline().split()]
        temp_split = file_id.readline().split()

        d = [float(temp_split[0])]
        for col in cols:
            d.append(float(temp_split[col]))
        data.append(d)

    file_id.close()
    axis = ['time']
    for i in range(len(cols)):
        axis.append('degree{}'.format(i))
    data = pd.DataFrame(data).set_axis(axis, axis=1)
    #data = data.set_index('time')
    return header, labels, data

def save_motion_file(data: pd.DataFrame, filename: str):
    with open(filename, 'w') as fo:
        fo.write("Coordinates\n")
        fo.write("version=1\n")
        fo.write("nRows={}\n".format(data.shape[0]))
        fo.write("nColumns={}\n".format(data.shape[1] - 1))
        fo.write("inDegrees=yes\n")
        fo.write("endheader\n")
        for k in data.columns:
            fo.write("{:>16}\t".format(k))
        fo.write("\n")
        for i in range(data.shape[0]):
            # fo.write("{:16.8f}\t".format(step_time * i))
            for col in data.columns:
                fo.write("{:16.8f}\t".format(data.loc[i, col]))
            fo.write("\n")
