# coding: utf-8

import os
import shutil
import subprocess
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.patches import Wedge
from scipy.integrate import trapz
from scipy.interpolate import NearestNDInterpolator, griddata, interp1d
from scipy.optimize import minimize

TESLAMAX_PACKAGE_DIR = Path(os.path.dirname(__file__))

TESLAMAX_JAVA_DIR = TESLAMAX_PACKAGE_DIR.parent / 'java'

TESLAMAX_CLASS_FILE = TESLAMAX_JAVA_DIR / 'TeslaMax.class'

TESLAMAX_CMD = ['comsolbatch', '-inputfile', str(TESLAMAX_CLASS_FILE)]

B_HIGH_FILENAME = "B_high.txt"
B_LOW_FILENAME = "B_low.txt"

B_III_FILENAME = "B_III.txt"

H_IV_FILENAME = "H_IV_1Q.txt"

MAIN_RESULTS_FILENAME = "COMSOL Main Results.txt"

MAGNETIC_PROFILE_FILENAME = "COMSOL Magnetic Profile.txt"

PARAMETER_FILENAME = "params.txt"

N_PROFILE_POINTS = 181  # keep at this level to have in increments of 1 degree
N_R_POINTS = 20

N_POINTS_PER_AXIS = 400

FIGSIZE_CM = 20
FIGSIZE_INCHES = FIGSIZE_CM / 2.54

FONTSIZE = 20

B_HIGH_LEVEL = 1.0
B_LOW_LEVEL = 0.0

DEBUG = False

def get_comsol_parameters_series(filename=PARAMETER_FILENAME):
    """Parse a COMSOL parameters file 'filename' and
    return a pandas Series from it.

    """
    param_comsol_file = Path(filename)

    param_comsol_series = pd.read_csv(str(param_comsol_file),
                                        squeeze=True,
                                        sep=" ",
                                        index_col=0,
                                        header=None)

    param_comsol_series.name = "COMSOL Parameters"
    param_comsol_series.index.name = None

    # append the units to the parameters names
    names_with_units = {}
    for name in param_comsol_series.keys():
        if name.startswith("h_") or name.startswith("R_"):
            names_with_units[name] = name + "[m]"
        if name.startswith("alpha") or name.startswith(
                "phi") or name.startswith("delta_phi"):
            names_with_units[name] = name + "[deg]"
        if name.startswith("B_"):
            names_with_units[name] = name + "[T]"
        if name.startswith("H_c"):
            names_with_units[name] = name + "[A/m]"

    param_comsol_series = param_comsol_series.rename(names_with_units)
    return param_comsol_series

def read_comsol_data_file(filename):
    """Read and parse 'filename' as exported by COMSOL.
    Export the numerical data as a numpy array containing only the numerical
    data; the first two columns are x and y values. All values are in SI.

    Keyword Arguments:
    filename -- str
    """

    return np.loadtxt(filename, skiprows=9)

def read_comsol_profile_data(filename):
    """
    Read 'filename' as exported by TeslaMax and return an array of the
    magnetic profile data, where the first column is the angle
    in degrees [0,360] and the second is the average magnetic
    flux density in tesla
    """

    profile_data = np.loadtxt(filename,
                              skiprows=1)

    return profile_data.T

def process_main_results_file():
    """Take the file "COMSOL Main Results.txt" as exported by COMSOL and
    clean the header data.

    """

    p = Path('.') / MAIN_RESULTS_FILENAME

    results = pd.read_csv(MAIN_RESULTS_FILENAME,
                            sep="\s+",
                            skiprows=5,
                            index_col=None,
                            header=None,
                            names=["B_high[T]",
                                   "B_low[T]",
                                   "A_gap[m2]",
                                   "A_magnet[m2]",
                                   "-H_Brem_II_max[A/m]",
                                   "-H_Brem_IV_max[A/m]",
                                   "A_demag[m2]"])

    results_series = results.ix[0]

    results_series.to_csv(str(p),
                          float_format="%.6f",
                          sep=" ",
                          header=False,
                          index=True)

def read_main_results_file():
    """Return a Series where each row is one of the COMSOL Main results"""

    results_filepath = Path(MAIN_RESULTS_FILENAME)

    results_series = pd.read_table(results_filepath,
                                   sep=" ",
                                   squeeze=True,
                                   index_col=0,
                                   header=None)
    results_series.index.name = None
    results_series.name = "COMSOL Main Results"
    return results_series

# noinspection PyPep8Naming
def calculate_magnitude(components_grid):
    """
    Return an array [x, y, norm(V)] from [x, y, Vx, Vy]
    """

    # noinspection PyPep8Naming
    x, y, Vx, Vy = components_grid.T

    V = np.sqrt(Vx * Vx + Vy * Vy)
    return np.array((x, y, V)).T #sem erro aqui mas sai um zero

def calculate_magnetic_profile(B_data, params):
    """
    Return the magnetic profile array [phi, B] based on data for the
    magnetic flux density [x, y, B] and a dictionary of parameters.

    The magnetic profile is defined as the magnetic flux density along the
    circumference in the middle of the air gap.

    The grid for 'B_data' is supposed to span the interval 0 <= phi <= 90
    (the first quadrant); this function mirrors this interval and return phi
    in the interval [0, 360].
    """

    params = expand_parameter_dictionary(params)

    R_g = params['R_g']
    R_o = params['R_o']

    # create ranges for phi and r
    phi_min = 0.0
    phi_max = np.pi / 2

    phi_vector_1q = np.linspace(phi_min, phi_max, N_PROFILE_POINTS)

    r_central = (R_o + R_g) / 2

    # calcualte the points (x,y) distributed along
    # radial lines
    x_grid = r_central * np.cos(phi_vector_1q)
    y_grid = r_central * np.sin(phi_vector_1q)

    B_profile_1q = griddata(B_data[:, 0:2],
                            B_data[:, 2],
                            np.array([x_grid, y_grid]).T)
    #ERRO AQUI. PRIMEIRO DIGITO DO B_profile_1q A PARTIR DE CERTO MOMENTO VEM NAN
    b=np.isnan(B_profile_1q)
    for i in range(0,len(b)):
        if b[i]==True:
            B_profile_1q[i]=B_profile_1q[i+1]
        else: continue

    # extrapolate data to the full circle
    phi_vector = np.concatenate((phi_vector_1q,
                                 phi_vector_1q + np.pi / 2,
                                 phi_vector_1q + np.pi,
                                 phi_vector_1q + (3 / 2) * np.pi))

    B_profile = np.concatenate((B_profile_1q,
                                B_profile_1q[::-1],
                                B_profile_1q,
                                B_profile_1q[::-1]))

    profile_data = np.array((np.rad2deg(phi_vector), B_profile)).T
    return profile_data

def write_magnetic_profile_file():
    """Create a file "COMSOL Magnetic Profile.txt" in the current directory,
    assuming the teslamax command was already ran, and write the magnetic
    profile data (magentic flux density at the air gap central circumference).
    """

    p = Path('.') / MAGNETIC_PROFILE_FILENAME
    column_names = ["phi[deg]", "B[T]"]
    column_header = " ".join(column_names)

    # load data from the B_III filename
    B_III_data = read_comsol_data_file(B_III_FILENAME)
    # get the columns corresponding to [x, y, B_x, B_y] and calculate [x, y, B]
    B_1q = calculate_magnitude(B_III_data[:, :4])

    case_series = get_comsol_parameters_series()

    profile_data = calculate_magnetic_profile(B_1q, case_series)

    np.savetxt(str(p),
               profile_data,
               fmt=("%.2f", "%.5f"),
               delimiter=" ",
               header=column_header,
               comments='')

def calculate_average_high_field(profile_data):
    """
        Return the average magnetic profile through the high field region,
        based on profile data [theta, B_profile]
        """

    # in the present model, the high field region (equivalent to the cold
    # blow in an AMR device) goes from -45° to +45°, and from 135° to 225°

    theta_min = 135.0
    theta_max = 225.0

    theta_vector = profile_data[:, 0]
    B_profile_vector = profile_data[:, 1]

    # select values only where theta_vector falls in the high field region
    # this will return an array with True only at those positions
    region_filter = np.logical_and(theta_vector > theta_min,
                                   theta_vector < theta_max)

    # select the region of the magnetic profile that satisfy this condition
    theta_high_field = theta_vector[region_filter]
    B_profile_high_field = B_profile_vector[region_filter]

    # return the integral of these samples, divided by the range
    B_integrated = trapz(B_profile_high_field, theta_high_field)
    theta_range = theta_max - theta_min

    B_high_avg = B_integrated / theta_range
    return B_high_avg

def write_magnetic_profile_central_file():
    """Create a file "COMSOL Magnetic Profile.txt" in the current directory,
    assuming the teslamax command was already ran, and write the magnetic
    profile data (magnetic induction at central radial position).

    """

    p = Path('.') / MAGNETIC_PROFILE_FILENAME
    column_names = ["phi[deg]", "B[T]"]
    column_header = " ".join(column_names)

    # load data from the high and low field regions
    B_h = read_comsol_data_file(B_HIGH_FILENAME)
    B_l = read_comsol_data_file(B_LOW_FILENAME)

    B_1q = np.concatenate((B_h, B_l), axis=0)

    # calcualte vector of angles for the first quadrant
    case_series = get_comsol_parameters_series()

    n_phi_points = 100

    R_g = case_series['R_g[m]']
    R_o = case_series['R_o[m]']

    # create ranges for phi and r
    phi_min = 0.0
    phi_max = np.pi / 2

    phi_vector_1q = np.linspace(phi_min, phi_max, N_PROFILE_POINTS)

    r_min = R_o
    r_max = R_g

    r_central = (R_o + R_g)/2

    # calcualte the points (x,y) distributed along
    # radial lines
    x_grid = r_central * np.cos(phi_vector_1q)
    y_grid = r_central * np.sin(phi_vector_1q)

    B_profile_1q = griddata(B_1q[:, 0:2], B_1q[:, 2],
                            np.array([x_grid, y_grid]).T)

    # extrapolate data to the full circle
    phi_vector = np.concatenate((phi_vector_1q,
                                 phi_vector_1q + np.pi / 2,
                                 phi_vector_1q + np.pi,
                                 phi_vector_1q + (3 / 2) * np.pi))

    B_profile = np.concatenate((B_profile_1q,
                                B_profile_1q[::-1],
                                B_profile_1q,
                                B_profile_1q[::-1]))

    profile_data = np.array((np.rad2deg(phi_vector), B_profile)).T

    np.savetxt(str(p),
               profile_data,
               fmt=("%.2f", "%.5f"),
               delimiter=" ",
               header=column_header,
               comments='')

def run_teslamax(verbose=False):
    """
    Run the teslamax process in the current directory, clean the results file
    and create a magnetic profile file.

    Assumes the parameters file is present in the current directory."""
    comsol_process = subprocess.run(TESLAMAX_CMD,
                                    shell=True,
                                    stdout=subprocess.PIPE,
                                    stderr=subprocess.STDOUT,
                                    universal_newlines=True)
    if verbose:
        print(comsol_process.stdout)
    process_main_results_file()
    write_magnetic_profile_file()

def remove_units_from_dict_keys(dictionary):
    """Remove a string '[<anything>]' from every key of 'dictionary'"""

    new_dictionary = {}

    for key in dictionary.keys():
        new_key = key.split('[')[0]
        new_dictionary[new_key] = dictionary[key]

    return new_dictionary

def expand_parameter_dictionary(param_simple):
    """
    Return a new dictionary, calculating derivative parameters from
    'param_simple', which is usually passed to COMSOL.

    If the input dictionary contain units, they are removed. The returned
    dict does not contain units
    """
    param_dict = remove_units_from_dict_keys(param_simple)

    # cast the number of segments to int, if necessary
    param_dict["n_IV"] = int(param_dict["n_IV"])

    # remove some keys that are not necessary (and which may cause errors)
    try:
        del param_dict["-H_Brem_IV_max"]
    except:
        pass

    # calculate magnet geometry
    param_dict["R_g"] = param_dict["R_o"] + param_dict["h_gap"]
    param_dict["R_c"] = param_dict["R_s"] + param_dict["h_fc"]
    if param_dict["n_II"] > 0:
        param_dict["delta_phi_S_II"] = ((param_dict["phi_S_II"] -
                                         param_dict["phi_C_II"]) /
                                        param_dict["n_II"])
    if param_dict["n_IV"] > 0:
        param_dict["delta_phi_S_IV"] = (param_dict["phi_S_IV"] /
                                        param_dict["n_IV"])

    return param_dict

def write_parameter_file_from_dict(param_dict):
    """From a basic 'param_dict', calculate the necessary other parameters
    (e.g. magnet segment size from total size and number of segments) and write
    the correct parameters file.

    If 'param_dict' contains units in the names, they are removed.
    """

    param_dict = expand_parameter_dictionary(param_dict)

    # write the dictionary file in the appropriate format that COMSOL can parse
    parameters_file_path = Path(".") / PARAMETER_FILENAME

    param_text = ""

    for (key, value) in param_dict.items():
        param_text = param_text + "%s %s\n" % (key, value)

    parameters_file_path.write_text(param_text)

def run_teslamax_from_params(params, verbose=False):
    """Write the 'params' dictionary in the apropriate format to the current
    directory (removing units if necessary) and run the teslamax process"""
    write_parameter_file_from_dict(params)
    run_teslamax(verbose)

def normalize_vector(v):
    """
    Return the normalized (dimensionless) form of vector
    (or list of vectors) v"""

    # v could be a single vector or a list of vectors,
    # so we handle different cases
    if v.ndim == 1:
        return v / np.linalg.norm(v)
    else:
        v_norm = np.linalg.norm(v, axis=1)
        v_norm_inv = np.reciprocal(v_norm).reshape(len(v), 1)
        return np.multiply(v, v_norm_inv)

def create_quater_circle_figure_template(r_lim, params):
    """
    Return (fig,axes) correspondent to a figure of the first quadrant,
    limited by r_lim.
    Both magnets are also drawn.

    The size of the figure is controlled by FIGSIZE_INCHES"""

    fig = plt.figure(figsize=(FIGSIZE_INCHES, FIGSIZE_INCHES))
    axes = fig.add_subplot(111, aspect='equal')

    axes.set_ylim(0, 1e3 * r_lim)
    axes.set_xlim(0, 1e3 * r_lim)

    axes.set_ylabel(r'$y\ [\si{\mm}$]')
    axes.set_xlabel(r'$x\ [\si{\mm}$]')

    R_o = params['R_o']
    R_i = params['R_i']
    R_s = params['R_s']
    R_g = params.get('R_g', params['R_o'] + params['h_gap'])

    magnet_II_outer = plt.Circle((0, 0), 1000 * R_o, color='k', fill=False)
    magnet_II_inner = plt.Circle((0, 0), 1e3 * R_i, color='k', fill=False)
    axes.add_artist(magnet_II_outer)
    axes.add_artist(magnet_II_inner)

    magnet_IV_outer = plt.Circle((0, 0), 1e3 * R_s, color='k', fill=False)
    magnet_IV_inner = plt.Circle((0, 0), 1000 * R_g, color='k', fill=False)
    axes.add_artist(magnet_IV_outer)
    axes.add_artist(magnet_IV_inner)

    return fig, axes

def generate_sector_mesh_points(R1, R2, phi1, phi2):
    """
    Return a list of points [X,Y] uniformily distributed in a circle between
    radii R1 and R2 and angular positions phi1 and phi2

    The number of points is controlled by N_POINTS_PER_AXIS.
    """

    phi_min = phi1
    phi_max = phi2

    phi_vector = np.linspace(phi_min, phi_max, N_POINTS_PER_AXIS)

    r_vector = np.linspace(R1, R2, N_POINTS_PER_AXIS)

    phi_grid, r_grid = np.meshgrid(phi_vector, r_vector)

    X_vector = (r_grid * np.cos(phi_grid)).flatten()
    Y_vector = (r_grid * np.sin(phi_grid)).flatten()
    return np.array([X_vector, Y_vector]).T

def create_magnet_IV_figure_template(params):
    """
    Return (fig,axes) correspondent to a figure of the
    first quadrant of magnet IV.

    The size of the figure is controlled by FIGSIZE_INCHES"""

    fig = plt.figure(figsize=(FIGSIZE_INCHES, FIGSIZE_INCHES))
    axes = fig.add_subplot(111, aspect='equal')

    R_o = params['R_o']
    R_i = params['R_i']
    R_s = params['R_s']
    R_g = params.get('R_g', params['R_o'] + params['h_gap'])
    R_c = params.get('R_c', params['R_s'] + params['h_fc'])
    r_lim = R_c

    axes.set_ylim(0, 1e3 * r_lim)
    axes.set_xlim(0, 1e3 * r_lim)

    axes.set_ylabel(r'$y\ [\si{\mm}$]')
    axes.set_xlabel(r'$x\ [\si{\mm}$]')

    width_IV = R_s - R_g
    n_IV = int(params['n_IV'])
    delta_phi_S_IV = params['delta_phi_S_IV']
    for i in range(0, n_IV):
        theta_0 = i * delta_phi_S_IV
        theta_1 = (i + 1) * delta_phi_S_IV
        magnet_segment = Wedge((0, 0),
                               1e3 * R_s,
                               theta_0,
                               theta_1,
                               1e3 * width_IV,
                               color='k',
                               fill=False)
        axes.add_artist(magnet_segment)

    return fig, axes

def create_magnets_figure_template(params):
    """
    Return (fig,axes) correspondent to a figure of the
    first quadrant of both magnets.

    The size of the figure is controlled by FIGSIZE_INCHES"""

    fig = plt.figure(figsize=(FIGSIZE_INCHES, FIGSIZE_INCHES))
    axes = fig.add_subplot(111, aspect='equal')

    R_o = params['R_o']
    R_i = params['R_i']
    R_s = params['R_s']
    R_g = params.get('R_g', params['R_o'] + params['h_gap'])
    R_c = params.get('R_c', params['R_s'] + params['h_fc'])
    r_lim = R_c

    axes.set_ylim(0, 1e3 * r_lim)
    axes.set_xlim(0, 1e3 * r_lim)

    axes.set_ylabel(r'$y\ [\si{\mm}$]')
    axes.set_xlabel(r'$x\ [\si{\mm}$]')

    width_II = R_o - R_i
    n_II = int(params['n_II'])
    delta_phi_S_II = params['delta_phi_S_II']
    for i in range(0, n_II):
        theta_0 = i * delta_phi_S_II
        theta_1 = (i + 1) * delta_phi_S_II
        magnet_segment = Wedge((0, 0),
                               1e3 * R_o,
                               theta_0,
                               theta_1,
                               1e3 * width_II,
                               color='k',
                               fill=False)
        axes.add_artist(magnet_segment)

    width_IV = R_s - R_g
    n_IV = int(params['n_IV'])
    delta_phi_S_IV = params['delta_phi_S_IV']
    for j in range(0, n_IV):
        theta_0 = j * delta_phi_S_IV
        theta_1 = (j + 1) * delta_phi_S_IV
        magnet_segment = Wedge((0, 0),
                               1e3 * R_s,
                               theta_0,
                               theta_1,
                               1e3 * width_IV,
                               color='k',
                               fill=False)
        axes.add_artist(magnet_segment)

    return fig, axes

class TeslaMaxGeometry:
    """
    Class representing the physical geometry of the TeslaMax system,
    with all radii and angles.

    To instantiante, pass a dictionary (or similar object) with all geometric
    parameters in SI units (except for the angles, which must be provided in
    degrees). The names for the keys follow the standard convention, without
    the units in the names. E.g.

    >>> params = {'R_i': 0.015, 'phi_C_II': 15, ...} # provide other parameters
    >>> tmg = TeslaMaxGeometry(params)

    The parameters 'R_o', 'h_gap' and 'R_g' are not independent. If two are
    provided, the class automatically calculates the other one. If you provide
    all three, it's your responsibility to provide three consistent values.

    Currently, the only possible calculations are volume-related.
    """

    def __init__(self, params):
        """
        Keyword Arguments:
        params -- dict-like
        """

        self.geometric_parameters = params.copy()
        self._complete_geometric_parameters()

    def _complete_geometric_parameters(self):
        """
        For two of the parameters 'R_o', 'R_g', 'h_gap', calculate the
        third one and populate the 'geometric_parameters' field.

        If all three parameters are provided, nothing happens
        """

        gp = self.geometric_parameters

        if ('R_o' in gp) and ('R_g' in gp) and ('h_gap' in gp):
            pass
        else:
            if ('R_o' in gp) and ('R_g' in gp):

                gp['h_gap'] = gp['R_g'] - gp['R_o']

            elif ('R_o' in gp) and ('h_gap' in gp):

                gp['R_g'] = gp['R_o'] + gp['h_gap']

            elif ('R_g' in gp) and ('h_gap' in gp):

                gp['R_o'] = gp['R_g'] - gp['h_gap']

    def calculate_magnet_volume(self, L):
        """
        Return the volume (m3) of the permanent regions,
        for a length of 'L' (m)
        """

        params = self.geometric_parameters

        phi_S_II = np.deg2rad(params["phi_S_II"])
        phi_S_IV = np.deg2rad(params["phi_S_IV"])
        phi_C_II = np.deg2rad(params["phi_C_II"])

        R_i = params["R_i"]
        R_o = params["R_o"]
        R_s = params["R_s"]
        R_g = params["R_g"]

        # the factor of 2 already accounts for 4 quadrants
        A_II = 2 * (phi_S_II - phi_C_II) * (R_o ** 2 - R_i ** 2)
        A_IV = 2 * phi_S_IV * (R_s ** 2 - R_g ** 2)
        V = (A_II + A_IV) * L

        return V

def expand_parameters_from_remanence_array(magnet_parameters, params, prefix):
    """
    Return a new parameters dict with the magnet parameters in the form
    '<prefix>_<magnet>_<segment>', with the values from 'magnet_parameters'
    and other parameters from 'params'.

    The length of the array 'magnet_parameters' must be equal to the sum of
    the number of segments in both cylinders.

    The first n_II elements refer to the inner magnet,
    and the remaining elements to the outer magnet.
    """

    params_expanded = params.copy()

    n_II = params["n_II"]
    for i in range(0, n_II):
        params_expanded["%s_II_%d" % (prefix, i + 1,)] = magnet_parameters[i]

    n_IV = params["n_IV"]
    for j in range(0, n_IV):
        k = j + n_II  # the first n_II elements refer to magnet II
        params_expanded["%s_IV_%d" % (prefix, j + 1,)] = magnet_parameters[k]

    return params_expanded

def calculate_instantaneous_profile(phi, B_high, B_low, *args):
    """
    Calculate the value of the two-pole instantaneous magnetic profile at
    angular position 'phi' (in degrees), where the profile oscillates from
    'B_low' to 'B_high'

    """

    high_region = (phi <= 45)
    high_region = np.logical_or(high_region,
                                np.logical_and((phi >= 135),
                                               (phi <= 225)))
    high_region = np.logical_or(high_region, (phi >= 315))
    return np.where(high_region, B_high, B_low)

def calculate_ramp_profile(phi, B_high, B_low, high_field_fraction, *args):
    """
    Calculate the value of the two-pole instantaneous magnetic profile at
    angular position 'phi' (in degrees), where the profile oscillates from
    'B_low' to 'B_high' in a trapezoidal wave, with each plateau occupying
    'high_field_fraction' of the cycle.

    """

    # for the edge case of a field fraction of 50%,
    # the ramp profile is equivalent to the instantaneous profile
    if np.isclose(high_field_fraction,0.5):
        return calculate_instantaneous_profile(phi,B_high,B_low,args)

    # for two poles, we can replicate the results from 0 to 180
    phi = np.mod(phi,180)

    # the fraction of the cycle where the field is constant is the fraction
    # where the field is at the high level, plus the fration where the field is
    # at the low level, hence the factor of 2
    field_fraction = 2 * high_field_fraction
    angle_change = field_fraction * 45

    high_region = (phi < angle_change)
    high_region = np.logical_or(high_region, (phi > (180 - angle_change)))

    descent_region = np.logical_and((phi >= angle_change),
                                    (phi <= 90 - angle_change))

    ascent_region = np.logical_and((phi >= 90 + angle_change),
                                   (phi <= 180 - angle_change))

    return np.where(high_region,
                    B_high,
                    np.where(descent_region,
                             B_high + (B_low - B_high) * (
                                     phi - angle_change) / (
                                         (1 - field_fraction) * 90),
                             np.where(ascent_region,
                                      B_low + (B_high - B_low) * (
                                              phi - (90 + angle_change)) / (
                                              (1 - field_fraction) * 90),
                                      B_low)))

class TeslaMaxPreDesign:
    """
    Class representing a fixed-geometry pre-design of the TeslaMax system,
    with all geometric and material parameters, but without the direction of
    magnetization for the magnet segments.

    To instantiate, you have to pass a dictionary with the parameters:
    >>> tmpd = TeslaMaxPreDesign({'R_i': 0.015, 'mu_r_II': 1.05, ...})

    This dictionary should at least contain all geometric parameters,
    with which a TeslaMaxGeometry object will be created. The material
    properties can be added  directly in the constructor;
    the remanences magnitudes for each segment in this case are specified
    as a vector:

    >>> geom_params = {'R_i': 0.015, 'n_II': 2, ...}
    >>> tmpd = TeslaMaxPreDesign(geom_params,
                                 mu_r_II=1.05,
                                 mu_r_IV=1.10,
                                 mu_r_iron=5e5,
                                 linear_iron=1,
                                 B_rem_vector=np.array([1.4,1.4,1.2,1.2])

    The parameters dictionary may contain some of the material properties;
    if you provide a parameter in the dictionary and via the constructor,
    it is the latter value that will be used to build the object. In the
    dictionary, the remanences must be provided in the usual form
    (e.g. 'B_rem_II_1')
    """

    def __init__(self,
                 params,
                 mu_r_II=None,
                 mu_r_IV=None,
                 B_rem_vector=None,
                 mu_r_iron=None,
                 linear_iron=None):

        self.geometry = TeslaMaxGeometry(params)
        self.geometry_material_parameters = self.geometry.geometric_parameters

        if mu_r_II is not None:
            self.geometry_material_parameters['mu_r_II'] = mu_r_II

        if mu_r_IV is not None:
            self.geometry_material_parameters['mu_r_IV'] = mu_r_IV

        if mu_r_iron is not None:
            self.geometry_material_parameters['mu_r_iron'] = mu_r_iron

        if linear_iron is not None:
            self.geometry_material_parameters['linear_iron'] = linear_iron

        if B_rem_vector is not None:
            self.geometry_material_parameters = expand_parameters_from_remanence_array(
                B_rem_vector,
                self.geometry_material_parameters,
                'B_rem')

        self.points_F_operators = None
        self.F_operators = None

        self.alpha_B_rem_optimal = None
        self.optimization_results = None

    def calculate_B_III_from_single_block(self,
                                          point,
                                          segment,
                                          magnet,
                                          magnitude,
                                          angle):
        """
        Return B_III(point) when 'segment' (1, 2, 3, ...)  of 'magnet'
        (either 'II' or 'IV') has a remanence of 'magnitude' and 'angle',
        and all other segments have null remanence.

        """

        n_II = self.geometry_material_parameters["n_II"]
        n_IV = self.geometry_material_parameters["n_IV"]
        n_total = n_II + n_IV

        B_rem_vector = np.zeros(n_total)
        alpha_B_rem_vector = np.zeros(n_total)

        if magnet == "II":
            element = segment - 1

        else:
            element = n_II + (segment - 1)

        B_rem_vector[element] = magnitude
        alpha_B_rem_vector[element] = angle

        tmpd = TeslaMaxPreDesign(self.geometry_material_parameters,
                                 B_rem_vector=B_rem_vector)

        # the results of these intermediate calculations are stored in this dir
        auxdir = Path('.') / 'teslamax-optimization'
        auxdir.mkdir(exist_ok=True)
        tmm = TeslaMaxModel(tmpd, alpha_B_rem_vector, str(auxdir))
        tmm.run(verbose=DEBUG)
        result = tmm.calculate_B_III_from_position(point)
        return result

    def calculate_F_operators(self):
        """
        Return (F_II_x, F_II_y, F_IV_x, F_IV_y), where each element is a list
        of the F-operators vector fields calculated at a uniform mesh
        in the air gap.

        For instance, F_II_x[0] is an array of (B_x, B_y), calculated when
        only the first segment of magnet II is magnetized in the x-direction,
        with unit remanence
        """

        n_II = self.geometry_material_parameters['n_II']
        n_IV = self.geometry_material_parameters['n_IV']

        R_o = self.geometry_material_parameters["R_o"]
        R_g = self.geometry_material_parameters["R_g"]
        points = generate_sector_mesh_points(1.001 * R_o,
                                             0.999 * R_g,
                                             0.0,
                                             np.pi / 2)
        self.points_F_operators = points

        F_II_x = []
        F_II_y = []

        F_IV_x = []
        F_IV_y = []

        for k in range(0, n_II):
            F_II_x.append(self.calculate_B_III_from_single_block(point=points,
                                                                 segment=k + 1,
                                                                 magnet='II',
                                                                 magnitude=1.0,
                                                                 angle=0.0))

            F_II_y.append(self.calculate_B_III_from_single_block(point=points,
                                                                 segment=k + 1,
                                                                 magnet='II',
                                                                 magnitude=1.0,
                                                                 angle=90.0))

        for j in range(0, n_IV):
            F_IV_x.append(self.calculate_B_III_from_single_block(point=points,
                                                                 segment=j + 1,
                                                                 magnet='IV',
                                                                 magnitude=1.0,
                                                                 angle=0.0))

            F_IV_y.append(self.calculate_B_III_from_single_block(point=points,
                                                                 segment=j + 1,
                                                                 magnet='IV',
                                                                 magnitude=1.0,
                                                                 angle=90.0))

        self.F_operators = (F_II_x, F_II_y, F_IV_x, F_IV_y)
    def get_points_F_operators(self):
        if self.points_F_operators is None:
            self.calculate_F_operators()
        return self.points_F_operators

    def get_F_operators(self):

        if self.F_operators is None:
            self.calculate_F_operators()
        return self.F_operators

    def superposition_B_III(self, alpha_B_rem):
        """
        Return (x, y, B_x, B_y) based on a  vector of remanence angles.

        - 'alpha_B_rem' is a vector of (n_II + n_IV) remanences, where the
        first n_II elements represent magnet II and the remaining elements
        represent magnet IV

        """

        B_III = 0.0

        points = self.get_points_F_operators()
        F_II_x, F_II_y, F_IV_x, F_IV_y = self.get_F_operators() #nao eh aqui

        params = self.geometry_material_parameters

        n_II = params["n_II"]
        for k in range(0, n_II):
            B_rem = params["B_rem_II_%d" % (k + 1)]
            alpha = np.deg2rad(alpha_B_rem[k])
            B = B_rem * (np.cos(alpha) * F_II_x[k] + np.sin(alpha) * F_II_y[k])
            B_III = B_III + B
        n_IV = params["n_IV"]
        for j in range(0, n_IV):
            B_rem = params["B_rem_IV_%d" % (j + 1)]
            alpha = np.deg2rad(alpha_B_rem[n_II + j])
            B = B_rem * (np.cos(alpha) * F_IV_x[j] + np.sin(alpha) * F_IV_y[j])

            B_III = B_III + B
        B_III_grid = np.concatenate((points, B_III), axis=1)

        return B_III_grid

    def calculate_functional_average(self, alpha_B_rem):
        """
        Return the objective functional based on  a vector of remanence angles.
        The objective functional is defined as the reciprocal of the
        average high field, to be minimized.

        - 'alpha_B_rem' is a vector of (n_II + n_IV) remanences, where the
        first n_II elements represent magnet II and the remaining elements
        represent magnet IV
        """

        B_III_data = self.superposition_B_III(alpha_B_rem)

        # the above statement will return [x,y,B_x,B_y]. We have to calculate the magnitude to pass it
        # to the magnetic profile data
        B_III_data = calculate_magnitude(B_III_data)

        B_profile_data = calculate_magnetic_profile(B_III_data,
                                                    self.geometry_material_parameters).T

        S = -calculate_average_high_field(B_profile_data)
        return S

    def calculate_functional_target(self,
                                    alpha_B_rem,
                                    target_profile_function,
                                    target_profile_args):
        """
        Return the objective functional based on  a vector of remanence angles.
        The objective functional is defined as the difference between the
        resulting profile and a target profile function,
        and is to be minimized.

        - 'alpha_B_rem' is a vector of (n_II + n_IV) remanences, where the
        first n_II elements represent magnet II and the remaining elements
        represent magnet IV
        - 'target_profile_function' is a function with signature
        'f(phi_vector, *args)' (the first argument is the vector of angular
        positions where the profile is to be calculated, followed by
        all other arguments)
        - 'target_profile_args' is a tuple with other arguments to pass to
        'target_profile_function' (see above). The first two elements
        are some measure of the maximum and minimum field
        """

        B_III_data = self.superposition_B_III(alpha_B_rem)

        # the above statement will return [x,y,B_x,B_y]. We have to calculate
        # the magnitude to pass it to the magnetic profile data
        B_III_data = calculate_magnitude(B_III_data)

        phi_vector, B_profile = calculate_magnetic_profile(B_III_data,
                                                           self.geometry_material_parameters).T

        B_target_profile = target_profile_function(phi_vector,
                                                   *target_profile_args)

        phi_vector = np.deg2rad(phi_vector)

        # use a "least squares" approach
        B_lsq = np.square((B_profile - B_target_profile))
        B_max = target_profile_args[0]
        B_min = target_profile_args[1]

        S = np.trapz(B_lsq, phi_vector) / (2 * np.pi * (B_max - B_min) ** 2)

        return S

    def calculate_functional(self,
                             alpha_B_rem,
                             functional_args=(calculate_instantaneous_profile,
                                              (B_HIGH_LEVEL,B_LOW_LEVEL))):
        """
        Return the objective functional based on  a vector of remanence angles.
        The objective functional is defined as the difference between the
        resulting profile and a target profile function,
        and is to be minimized.

        - 'alpha_B_rem' is a vector of (n_II + n_IV) remanences, where the
        first n_II elements represent magnet II and the remaining elements
        represent magnet IV
        - `functional_args' is a tuple in the form
        (target_profile_function, target_profile_args), where:
            - 'target_profile_function' is a function with signature
        'f(phi_vector, *args)' (the first argument is the vector of angular
        positions where the profile is to be calculated, followed by
        all other arguments)
            - 'target_profile_args' is a tuple with other arguments to pass to
        'target_profile_function' (see above)
        """

        return self.calculate_functional_target(alpha_B_rem,
                                                functional_args[0],
                                                functional_args[1])

    def calculate_functional_derivative(self,
                                        alpha_B_rem,
                                        i,
                                        functional_args):
        """
        Return the derivative of the functional in respect to the i-th element
        of the remanence angles vector.

        - 'alpha_B_rem' is a vector of (n_II + n_IV) remanences, where the
        first n_II elements represent magnet II and the remaining elements
        represent magnet IV
        - 'i' is the element (0-based) in respect to which the derivative
        is being evaluated
        - 'functional_args' is a tuple with other arguments that are passed to
        the functional method
        """

        S = self.calculate_functional(alpha_B_rem,
                                      functional_args)

        alpha_B_rem_plus = alpha_B_rem.copy()
        delta = 1e-6
        alpha_B_rem_plus[i] = alpha_B_rem_plus[i] + delta

        S_plus = self.calculate_functional(alpha_B_rem_plus,
                                           functional_args)

        dS = (S_plus - S) / delta
        return dS

    def calculate_funcional_derivative_second_order(self,
                                                    alpha_B_rem,
                                                    i,
                                                    j,
                                                    functional_args):
        """
        Return the second-order derivative of the functional in respect
        to the (i,j) elements (e.g. d/dalpha_i (dfunctional/dalpha_j))
        """

        dS_j = self.calculate_functional_derivative(alpha_B_rem,
                                                    j,
                                                    functional_args)

        alpha_B_rem_plus = alpha_B_rem.copy()
        delta = 1e-6
        alpha_B_rem_plus[i] = alpha_B_rem_plus[i] + delta

        dS_j_plus = self.calculate_functional_derivative(alpha_B_rem_plus,
                                                         j,
                                                         functional_args)

        ddS = (dS_j_plus - dS_j) / delta
        return ddS

    def calculate_functional_gradient(self,
                                      alpha_B_rem,
                                      functional_args=(
                                              calculate_instantaneous_profile,
                                              (B_HIGH_LEVEL,))):
        """
        Return the gradient of the functional evaluated at point 'alpha_B_rem'.

        Arguments:
        - alpha_B_rem is a vector of (n_II + n_IV) remanences, where the
        gradient is to be evaluated
        - 'functional_args' is a tuple with other arguments that are passed to
        the functional method
        """

        n = len(alpha_B_rem)

        grad = np.array([self.calculate_functional_derivative(alpha_B_rem,
                                                              i,
                                                              functional_args)
                         for i in range(0, n)])
        return grad

    def calculate_functional_hessian(self,
                                     alpha_B_rem,
                                     functional_args=(
                                             calculate_instantaneous_profile,
                                             (B_HIGH_LEVEL,))):
        """
        Return the Hessian matrix of the functional evaluated at
        point 'alpha_B_rem'.

        Arguments:
        - alpha_B_rem is a vector of (n_II + n_IV) remanences, where the
        gradient is to be evaluated
        - 'functional_args' is a tuple with other arguments that are passed to
        the functional method
        """

        n = len(alpha_B_rem)

        hess = np.array([
            [self.calculate_funcional_derivative_second_order(alpha_B_rem,
                                                              i,
                                                              j,
                                                              functional_args)
             for j in range(0, n)]
            for i in range(0, n)])
        return hess

    def _calculate_optimal_remanence_angles(self,
                                            target_profile_function=calculate_instantaneous_profile,
                                            target_profile_args=(
                                                    B_HIGH_LEVEL,)):

        """
        Calculate the optimal remanence angles that minimize the deviation
        between the resulting profile and 'target_profile_function'.

        Arguments:
        ----------

        - 'target_profile_function' is a function with signature
        'f(phi_vector, *args)' (the first argument is the vector of angular
        positions where the profile is to be calculated, followed by
        all other arguments)
        - 'target_profile_args' is a tuple with other arguments to pass to
        'target_profile_function' (see above)

        --
        """

        n_II = self.geometry_material_parameters["n_II"]
        n_IV = self.geometry_material_parameters["n_IV"]

        n = n_II + n_IV

        alpha_B_rem_0 = np.zeros(n)

        # this function, to be used by the minimize function, has signature
        # f(alpha_B_rem,args); i.e. it takes a candidate vector of
        # remanence angles and a tuple of other arguments
        objective_function = self.calculate_functional

        # in the case of our functional formulation, the above function
        # takes the name of the target profile and other parameters that are
        # passed along
        functional_args = (target_profile_function, target_profile_args)

        bounds = [(0.0, 360.0) for i in range(0, n)]

        optres_g = minimize(objective_function,
                            alpha_B_rem_0,
                            args=(functional_args,),
                            bounds=bounds,
                            options={'disp': False})

        self.optimization_results = optres_g
        self.alpha_B_rem_optimal = optres_g.x

    def get_optimal_remanence_angles(self,
                                     target_profile_function=calculate_instantaneous_profile,
                                     target_profile_args=(B_HIGH_LEVEL,)):

        """
        Return the optimal remanence angles that minimize the deviation
        between the resulting profile and 'target_profile_function'.

        Arguments:
        ----------

        - 'target_profile_function' is a function with signature
        'f(phi_vector, *args)' (the first argument is the vector of angular
        positions where the profile is to be calculated, followed by
        all other arguments)
        - 'target_profile_args' is a tuple with other arguments to pass to
        'target_profile_function' (see above)

        """

        if self.alpha_B_rem_optimal is None:
            self._calculate_optimal_remanence_angles(target_profile_function,
                                                     target_profile_args)

        return self.alpha_B_rem_optimal

class TeslaMaxModel:
    """
    Class representing the full TeslaMax model.

    To create an instance of the class, first you have to create a
    TeslaMaxPreDesign object and an array of remanence angles,
    and then pass these parameters along with  a path where to store the
    generated text files:

    >>> tmpd = TeslaMaxPreDesign({...})
    >>> alpha_B_rem = [...]
    >>> tmm = TeslaMaxModel(tmpd, alpha_B_rem, "teslamax-results")

    If the path already exists, it will be cleaned up to avoid confusion
    between different simulations.

    If you want to preserve the data in 'path', set the argument 'clean'
    to False in the constructor

    """

    def __init__(self, teslamax_predesign, alpha_vector, path, clean=True):
        self.pre_design = teslamax_predesign

        self.params = expand_parameters_from_remanence_array(alpha_vector,
                                                             self.pre_design.geometry_material_parameters,
                                                             "alpha_rem")

        self.path = Path(path)
        if (self.path.exists()) and (self.path.is_dir()) and clean:
            shutil.rmtree(str(self.path))

        self.path.mkdir(exist_ok=True)

        self.profile_data = np.empty(2)
        self.calculate_B_profile = None
        self.B_III_data = None
        self.calculate_B_III_from_position = None
        self.demagnetized_volume_fraction = None

    def run(self, verbose=False):
        """
        Change the current directory temporarily to the object's path
        and run the TeslaMax program in it.

        Also populates the appropriate fields with results from TeslaMax
        """

        cwd = os.getcwd()
        os.chdir(str(self.path))
        run_teslamax_from_params(self.params, verbose)
        os.chdir(cwd)

        self.profile_data = self.get_profile_data()
        self.calculate_B_profile = interp1d(self.profile_data[:, 0],
                                            self.profile_data[:, 1],
                                            kind='linear')

        self.B_III_data = self.get_B_III_data()
        self.calculate_B_III_from_position = NearestNDInterpolator(
            self.B_III_data[:, :2],
            self.B_III_data[:, 2:4])

        self.demagnetized_volume_fraction = self.get_results_series()["Demagnetized fraction[%]"]

    def get_B_III_data(self):
        """
        Return an array [x, y, Bx, By] for the air gap region
        """

        B_III_file_path = self.path / B_III_FILENAME
        B_III_full_data = read_comsol_data_file(str(B_III_file_path))

        B_III_data = B_III_full_data[:, :4]
        return B_III_data

    def get_profile_data(self):
        """
        Return an array of the magnetic profile data, where the first column
        is the angle in degrees [0,360] and the second is the average magnetic
        flux density in tesla
        """

        profile_file_path = self.path / MAGNETIC_PROFILE_FILENAME
        profile_data = read_comsol_profile_data(str(profile_file_path))

        return profile_data

    def get_results_series(self):
        """
        Return a Series with the main simulation results
        """

        results_filepath = self.path / MAIN_RESULTS_FILENAME

        results_series = pd.read_csv(results_filepath,
                                       sep=" ",
                                       squeeze=True,
                                       index_col=0,
                                       header=None)
        results_series.index.name = None
        results_series.name = "COMSOL Main Results"

        # add demagnetization fraction
        results_series["Demagnetized fraction[%]"] = results_series["A_demag[m2]"] / results_series["A_magnet[m2]"] * 100
        return results_series


# This module includes the functions that perform the "pre-processing": In short, they prepare the data structure
# that is then processed by the "Energy analysis" and "Exergy analysis" functions, all in one go.
#
# The module is made of :
# - One function that simply reads in the data that is already available from the original dataset

import numpy as np
import pandas as pd

def mainEngineProcessing(raw, processed, CONSTANTS, status, hd):
    # This script summarizes all the functions that calculate the required data for the Main Engines different flows
    # Reading existing values
    processed = readMainEnginesExistingValues(raw, processed, CONSTANTS, hd)
    # Calculating the main engines fuel flows
    processed = mainEngineFuelFlowCalculation(raw, processed, CONSTANTS, hd)
    # Calculating the main engines power output
    processed = mainEnginePowerCalculation(processed, CONSTANTS)
    # Calculating engine load, that is used many times later on
    status = engineStatusCalculation("MainEngines", raw, processed, CONSTANTS, status, hd)
    # Calculating air and exhaust gas flows in the main engines
    processed = mainEngineAirFlowCalculation(raw, processed, status, CONSTANTS)
    # Calculating cooling flows
    processed = engineCoolingSystemsCalculation(processed, CONSTANTS, status, "MainEngines")
    return (processed, status)

def auxEngineProcessing(raw, processed, CONSTANTS, status, hd):
    # This script summarizes all the functions that calculate the required data for the Main Engines different flows
    # Reading existing values
    processed = readAuxEnginesExistingValues(raw, processed, CONSTANTS, hd)
    # Calculating the power, including the generator efficiency
    processed = auxEnginePowerCalculation(processed, CONSTANTS)
    # Calculating engine load, that is used many times later on
    status = engineStatusCalculation("AuxEngines", raw, processed, CONSTANTS, status, hd)
    # Calculating the auxiliary engines fuel flows
    processed = auxEngineFuelFlowCalculation(raw, processed, CONSTANTS, status)
    # Calculate air and exhaust gas flows in the main engines
    processed = auxEngineAirFlowCalculation(raw, processed, CONSTANTS)
    # Calculating cooling flows
    processed = engineCoolingSystemsCalculation(processed, CONSTANTS, status, "AuxEngines", )
    return (processed, status)

def assumptions(raw, processed, CONSTANTS, hd):
    # This function includes generic assumed values in the main structure
    for name in CONSTANTS["General"]["NAMES"]["MainEngines"]:
        # The pressure at the turbocharger inlet is equal to the atmospheric pressure
        processed[name]["TC"]["Air_in"]["p"] = CONSTANTS["General"]["P_ATM"]
    for name in CONSTANTS["General"]["NAMES"]["AuxEngines"]:
        # The pressure at the turbocharger inlet is equal to the atmospheric pressure
        processed[name]["TC"]["Air_in"]["p"] = CONSTANTS["General"]["P_ATM"]
    processed["T_0"] = raw[hd["ER13_SW_T_IN"]] + 273.15
    return processed

def readMainEnginesExistingValues(raw, processed, CONSTANTS, hd):
    # This function only reads existing series. It does not do any pre-processing action.
    for name in CONSTANTS["General"]["NAMES"]["MainEngines"]:
        # Reading main engines exhaust gas temperature, TC inlet and outlet
        processed[name]["Cyl"]["EG_out"]["T"] = raw[hd[name + "-TC_EG_T_IN"]] + 273.15  # Measured before mixer with flow form bypass
        processed[name]["TC"]["EG_out"]["T"] = raw[hd[name + "-TC_EG_T_OUT"]] + 273.15  # Measured after mixer with waste gate
        # Reading main engines exhaust gas temperature, after HRSG. Only two of the four main engines have the HRSG
        if name=="ME2" or name=="ME3":
            processed[name]["HRSG"]["EG_out"]["T"] = raw[hd[name + "-EGB_EG_T_OUT"]] + 273.15
            processed[name]["HRSG"]["EG_in"]["T"] = raw[hd[name + "-TC_EG_T_OUT"]] + 273.15
        # Temperature in the engine room, i.e. inlet to the compressor of the TC
        processed[name]["TC"]["Air_in"]["T"] = raw[hd["ER-FWD_AIR_T_"]] + 273.15
        processed[name]["Comp"]["Air_in"]["T"] = processed[name]["TC"]["Air_in"]["T"]
        # Pressure of the charge air, at the compressor outlet (and, hence, at the cylinder inlet)
        processed[name]["TC"]["Air_out"]["p"] = raw[hd[name+"-CAC_AIR_P_OUT"]] + 1
        processed[name]["Cyl"]["Air_in"]["p"] = processed[name]["TC"]["Air_out"]["p"]
        processed[name]["BPvalve"]["Air_in"]["p"] = processed[name]["TC"]["Air_out"]["p"]
        # Reading the HT temperature before and after the main engine
        # processed[name]["CAC_HT"]["Water_out"]["T"] = raw[name + "_HT_water_T_out"] # Note: this might be inconsistent
        processed[name]["JWC"]["HTWater_in"]["T"] = raw[hd[name + "-HT_FW_T_IN"]] + 273.15
        # Reading the LT temperature before the main engine
        processed[name]["CAC_LT"]["LTWater_in"]["T"] = raw[hd[name + "-LT_FW_T_IN"]] + 273.15
        # Reading the Lubricating oil temperature before and after the Lubricating Oil Cooler (hence, In is higher)
        processed[name]["LOC"]["LubOil_out"]["T"] = raw[hd[name + "-LOC_OIL_T_OUT"]] + 273.15
        #                           processed[name]["LOC"]["LubOil_out"]["T"] = raw[hd[name + "_TC_OIL_T_OUT"]]
        # Reading fuel oil temperature before injection
        processed[name]["Cyl"]["FuelPh_in"]["T"] = raw[hd[name + "-CYL_FUEL_T_IN"]] + 273.15
        # Reading charge air temperature, after the charge air cooler (or at cylinder inlet)
        processed[name]["CAC_LT"]["Air_out"]["T"] = raw[hd[name + "-CAC_AIR_T_OUT"]] + 273.15
        processed[name]["Cyl"]["Air_in"]["T"] = processed[name]["CAC_LT"]["Air_out"]["T"]
        # Reading Engine rpm
        processed[name]["Cyl"]["Power_out"]["omega"] = raw[hd[name + "__RPM_"]]
    return processed

def mainEngineFuelFlowCalculation(raw, processed, CONSTANTS, hd):
    # This function calculates the engine
    for name in CONSTANTS["General"]["NAMES"]["MainEngines"]:
        # This function calculates the fuel flow of the main engines
        # In the case of the main engines, the fuel flow of an engine is calculated given its fuel
        # rack position and its rotating speed.
        fuel_rack_position = CONSTANTS["MainEngines"]["FRP_2_MFR"]["FRP_MIN"][name] + (CONSTANTS["MainEngines"]["FRP_2_MFR"]["FRP_MAX"][name]-CONSTANTS["MainEngines"]["FRP_2_MFR"]["FRP_MIN"][name]) * raw[hd[name+"__FRP_"]]/100
        # Temporarily, only the ISO fuel flow is calculated
        processed[name]["Cyl"]["FuelPh_in"]["mdot"] = CONSTANTS["MainEngines"]["MFR_FUEL_DES_ISO"] * (
            (CONSTANTS["MainEngines"]["FRP_2_MFR"]["POLY"][name][0] + CONSTANTS["MainEngines"]["FRP_2_MFR"]["POLY"][name][1] * fuel_rack_position) /
            (CONSTANTS["MainEngines"]["FRP_2_MFR"]["POLY"][name][0] + CONSTANTS["MainEngines"]["FRP_2_MFR"]["POLY"][name][1] * CONSTANTS["MainEngines"]["FRP_DES"][name])) * (
            processed[name]["Cyl"]["Power_out"]["omega"] / CONSTANTS["MainEngines"]["RPM_DES"])
        aaa = 0
    return processed

def mainEnginePowerCalculation(processed, CONSTANTS):
    # This function calculates the Power of the engine starting from the efficiency of the engine,
    # which is calcualted starting from other available data
    for name in CONSTANTS["General"]["NAMES"]["MainEngines"]:
        # Calculate fuel flow-based engine load
        fuel_based_load = processed[name]["Cyl"]["FuelPh_in"]["mdot"] / CONSTANTS["MainEngines"]["MFR_FUEL_DES_ISO"]
        # Calculate ISO bsfc (break specific fuel consumption)
        bsfc_iso = fuel_based_load.apply(polyvalHelperFunction, args=(CONSTANTS["MainEngines"][                                                                  "POLY_FUEL_LOAD_2_BSFC_ISO"],))
        # Corrects the bsfc from ISO conditions to "real" conditions
        (bsfc,LHV) = bsfcISOCorrection(bsfc_iso,processed[name]["Cyl"]["Air_in"]["T"],processed[name]["CAC_LT"]["LTWater_in"]["T"],processed[name]["Cyl"]["FuelPh_in"]["T"],CONSTANTS)
        # Calculates the real fuel flow using the ISO conversion
        processed[name]["Cyl"]["FuelPh_in"]["mdot"] = processed[name]["Cyl"]["FuelPh_in"]["mdot"] * bsfc / bsfc_iso
        # Calculates the power of the engine as mfr/bsfc, with unit conversion to get the output in kW
        # Shaft energy out
        processed[name]["Cyl"]["Power_out"]["Wdot"] = processed[name]["Cyl"]["FuelPh_in"]["mdot"] / bsfc * 1000 * 3600
        # Chemical energy in the fuel
        processed[name]["Cyl"]["FuelCh_in"]["Wdot"] = processed[name]["Cyl"]["FuelPh_in"]["mdot"] * LHV
    return processed

def mainEngineAirFlowCalculation(raw, processed, status, CONSTANTS):
    # This function calculates the different air and exhaust gas flows in the main engines, taking into account the
    # presence of air bypass and exhaust wastegate valves
    for name in CONSTANTS["General"]["NAMES"]["MainEngines"]:
        # Reading the pressure of the charge air after the compressor from the raw database
        # Calculating the compressor isentropic efficiency
        comp_isentropic_efficiency = processed[name]["Cyl"]["Air_in"]["p"].apply(polyvalHelperFunction,args=(CONSTANTS["MainEngines"][
                                                                                "POLY_PIN_2_ETA_IS"],))
        # Calculating the compressor's compression ratio
        beta_comp = processed[name]["TC"]["Air_out"]["p"] / processed[name]["TC"]["Air_in"]["p"]
        # Calculating the temperature after the compressor, based on ideal gas assumption
        processed[name]["TC"]["Air_out"]["T"] = processed[name]["TC"]["Air_in"]["T"] * beta_comp**((
            CONSTANTS["General"]["K_AIR"]-1)/CONSTANTS["General"]["K_AIR"]) / comp_isentropic_efficiency
        # The state of the air at the outlet of the compressor, outlet of the TC block, and inlet of the bypass valve
        #  are the same
        processed[name]["Comp"]["Air_out"]["T"] = processed[name]["TC"]["Air_out"]["T"]
        processed[name]["BPvalve"]["Air_in"]["T"] = processed[name]["TC"]["Air_out"]["T"]
        # The same temperature is the inlet to the charge air cooler
        processed[name]["CAC_HT"]["Air_in"]["T"] = processed[name]["TC"]["Air_out"]["T"]
        # Calculating the air inflow aspired by the cylinder: calculated as inlet air density times the maximum volume,
        # times the engine speed
        processed[name]["Cyl"]["Air_in"]["mdot"] = CONSTANTS["MainEngines"]["V_SW"] * (
            processed[name]["Cyl"]["Air_in"]["p"] * 1e5) / (
            CONSTANTS["General"]["R_AIR"] * processed[name]["Cyl"]["Air_in"]["T"]) * (
            processed[name]["Cyl"]["Power_out"]["omega"] / 60 / 2 * CONSTANTS["General"]["ETA_VOL"]) * (
            CONSTANTS["MainEngines"]["N_CYL"])
        processed[name]["Cyl"]["EG_out"]["mdot"] = processed[name]["Cyl"]["Air_in"]["mdot"] + processed[name]["Cyl"][
            "FuelPh_in"]["mdot"]
        # Here we calculate the bypass mass flow. THIS NEEDS TO BE CHECKED FOR CONSISTENCY (i.e. the bypass mass flow
        #  should always be kind of low, and anyway only be seen at low to medium load).
        # The equation is the result of a mass and energy balance over the whole engine
        #dh_comp_on_eta = (CONSTANTS["General"]["CP_AIR"] / CONSTANTS["MainEngines"]["ETA_MECH_TC"] *
        #    processed[name]["Cyl"]["Air_in"]["mdot"] * (
        #    processed[name]["Comp"]["Air_out"]["T"] - processed[name]["Comp"]["Air_in"]["T"]))
        #processed[name]["BPvalve"]["Air_in"]["mdot"] = (
        #    processed[name]["Cyl"]["EG_out"]["mdot"] * CONSTANTS["General"]["CP_EG"] * (
        #       processed[name]["Cyl"]["EG_out"]["T"] - processed[name]["TC"]["EG_out"]["T"]) - dh_comp_on_eta) / (
        #  CONSTANTS["General"]["CP_EG"] * (
        #   processed[name]["TC"]["EG_out"]["T"] - processed[name]["Comp"]["Air_out"]["T"]) -
        #   dh_comp_on_eta)
        #processed[name]["BPvalve"]["Air_in"]["mdot"] = (
        #    CONSTANTS["General"]["CP_AIR"] * processed[name]["Cyl"]["Air_in"]["mdot"] * (processed[name]["Comp"]["Air_out"]["T"] - processed[name]["Comp"]["Air_in"]["T"]) +
        #    CONSTANTS["General"]["CP_EG"] * CONSTANTS["MainEngines"]["ETA_MECH_TC"] * (processed[name]["Cyl"]["Air_in"]["mdot"] - processed[name]["Cyl"]["FuelPh_in"]["mdot"]) *
        #    (processed[name]["Cyl"]["EG_out"]["T"] - processed[name]["TC"]["EG_out"]["T"])) / (
        #    CONSTANTS["General"]["CP_AIR"] * (
        #        processed[name]["Comp"]["Air_in"]["T"]+ CONSTANTS["MainEngines"]["ETA_MECH_TC"] * processed[name]["TC"]["EG_out"]["T"] -
        #        (1 + CONSTANTS["MainEngines"]["ETA_MECH_TC"]) * processed[name]["Comp"]["Air_out"]["T"]))
        # The new approximation is that the valve is only open for engine load below 50%, and when it is open it increases the flow by a fixed amount
        processed[name]["BPvalve"]["Air_in"]["mdot"][:] = 0
        processed[name]["BPvalve"]["Air_in"]["mdot"][(status[name]["Load"]<0.5) | ((status[name]["Load"]<0.6) & (processed[name]["TC"]["EG_out"]["T"]<620))] = CONSTANTS["MainEngines"]["BYPASS_FLOW"]
        processed[name]["BPvalve"]["Air_out"]["mdot"] = processed[name]["BPvalve"]["Air_in"]["mdot"]
        # Calculating the temperature of the mixture after the merge between bypass and exhaust gas from the cylinders
        processed[name]["TC"]["EG_in"]["T"] = (
            processed[name]["BPvalve"]["Air_in"]["mdot"] * CONSTANTS["General"]["CP_AIR"] * processed[name]["Comp"]["Air_out"]["T"] +
            processed[name]["Cyl"]["Air_in"]["mdot"] * CONSTANTS["General"]["CP_EG"] * processed[name]["Cyl"]["EG_out"]["T"]) / (
            processed[name]["Cyl"]["EG_out"]["mdot"] * CONSTANTS["General"]["CP_EG"] +
            processed[name]["BPvalve"]["Air_in"]["mdot"] * CONSTANTS["General"]["CP_AIR"])
        # The air mass flow going through the compressor is equal to the sum of the air flow through the bypass valve and
        # to the cylinders
        processed[name]["TC"]["Air_in"]["mdot"] = processed[name]["BPvalve"]["Air_in"]["mdot"] + processed[name]["Cyl"]["Air_in"]["mdot"]
        # Inlet and outlet flows to the compressor are equal
        processed[name]["TC"]["Air_out"]["mdot"] = processed[name]["TC"]["Air_in"]["mdot"]
        # The flow through the turbine is equal to the sum of the bypass flow and the exhaust coming from the cylinders
        processed[name]["TC"]["EG_in"]["mdot"] = processed[name]["BPvalve"]["Air_in"]["mdot"] + processed[name]["Cyl"]["EG_out"]["mdot"]
        processed[name]["TC"]["EG_out"]["mdot"] = processed[name]["TC"]["EG_in"]["mdot"]
    return processed

def readAuxEnginesExistingValues(raw, processed,CONSTANTS,hd):
    for name in CONSTANTS["General"]["NAMES"]["AuxEngines"]:
        # Reading main engines exhaust gas temperature, TC inlet and outlet
        processed[name]["TC"]["EG_in"]["T"] = raw[hd[name + "-TC_EG_T_IN1"]] + 273.15
        processed[name]["TC"]["EG_out"]["T"] = raw[hd[name + "-TC_EG_T_OUT"]] + 273.15
        # Reading main engines exhaust gas temperature, after HRSG
        processed[name]["HRSG"]["EG_out"]["T"] = raw[hd[name + "-EGB_EG_T_OUT"]] + 273.15
        processed[name]["HRSG"]["EG_in"]["T"] = processed[name]["TC"]["EG_out"]["T"] + 273.15
        # Temperature in the engine room, i.e. inlet to the compressor of the TC
        processed[name]["TC"]["Air_in"]["T"] = raw[hd["ER_AIR_T_"]] + 273.15
        # Reading the HT temperature before and after the main engine
        processed[name]["JWC"]["HTWater_in"]["T"] = raw[hd[name + "-HT_FW_T_IN"]] + 273.15
        # Reading the LT temperature before the main engine
        processed[name]["CAC_LT"]["LTWater_in"]["T"] = raw[hd[name + "-LT-CAC_FW_T_IN"]] + 273.15
        # Reading the Lubricating oil temperature before and after the Lubricating oil cooler
        processed[name]["LOC"]["LubOil_out"]["T"] = raw[hd[name + "-LOC_OIL_T_OUT"]] + 273.15
        # Reading fuel oil temperature before injection
        processed[name]["Cyl"]["FuelPh_in"]["T"] = raw[hd[name + "-CYL_FUEL_T_IN"]] + 273.15
        # Reading charge air temperature.
        processed[name]["CAC_LT"]["Air_out"]["T"] = raw[hd[name + "-CAC_AIR_T_OUT"]] + 273.15
        processed[name]["CAC_LT"]["Air_out"]["p"] = raw[hd[name + "-CAC_AIR_P_OUT"]] + 1
        processed[name]["AG"]["Power_out"]["Wdot"] = raw[hd[name + "_POWER_Wdot_OUT"]]
        processed[name]["Cyl"]["Air_in"]["T"] = processed[name]["CAC_LT"]["Air_out"]["T"]
        processed[name]["Cyl"]["Air_in"]["p"] = processed[name]["CAC_LT"]["Air_out"]["p"]
    return processed

def auxEnginePowerCalculation(processed, CONSTANTS):
    # Calculating the power of the auxiliary engines
    for name in CONSTANTS["General"]["NAMES"]["AuxEngines"]:
        load = processed[name]["AG"]["Power_out"]["Wdot"] / CONSTANTS["AuxiliaryEngines"]["MCR"]
        eta_AG =  CONSTANTS["AuxiliaryEngines"]["AG"]["ETA_DES"] - CONSTANTS["AuxiliaryEngines"]["AG"]["A"] * np.exp(
            -CONSTANTS["AuxiliaryEngines"]["AG"]["k"] * (load))
        processed[name]["AG"]["Power_in"]["Wdot"] = processed[name]["AG"]["Power_out"]["Wdot"] / eta_AG
        processed[name]["AG"]["Losses"]["Wdot"] = processed[name]["AG"]["Power_in"]["Wdot"] - processed[name]["AG"]["Power_out"]["Wdot"]
        processed[name]["Cyl"]["Power_out"]["Wdot"] = processed[name]["AG"]["Power_in"]["Wdot"]

def auxEngineFuelFlowCalculation(raw, processed, CONSTANTS, status):
    # Proceeding with the auxiliary engines
    for name in CONSTANTS["General"]["NAMES"]["AuxEngines"]:
        # First we calculate the ISO break specific fuel consumption (BSFC)
        bsfc_iso = status[name]["Load"].apply(polyvalHelperFunction, args=(CONSTANTS["AuxEngines"][
                                                                          "POLY_LOAD_2_ISO_BSFC"],))
        (bsfc, LHV) = bsfcISOCorrection(bsfc_iso, processed[name]["Cyl"]["Air_in"]["T"], processed[name]["CAC_LT"][
                "LTWater_in"]["T"], processed[name]["Cyl"]["FuelPh_in"]["T"], CONSTANTS)
        processed[name]["Cyl"]["FuelPh_in"]["mdot"] = bsfc * processed[name]["Cyl"]["Power_out"]["Wdot"] / 3600 / 1000
        processed[name]["Cyl"]["FuelCh_in"]["Wdot"] = processed[name]["Cyl"]["FuelPh_in"]["mdot"] * LHV
    return processed

def auxEngineAirFlowCalculation(raw, processed, CONSTANTS):
    # This function calculates the different air and exhaust gas flows in the main engines, taking into account the
    # presence of air bypass and exhaust wastegate valves
    for name in CONSTANTS["General"]["NAMES"]["AuxEngines"]:
        # Reading the pressure of the charge air after the compressor from the raw database
        # Calculating the compressor isentropic efficiency
        comp_isentropic_efficiency = processed[name]["Cyl"]["Air_in"]["p"].apply(polyvalHelperFunction,args=(CONSTANTS["AuxEngines"][
                                                                                "POLY_PIN_2_ETA_IS"],))
        # Calculating the compressor's compression ratio
        beta_comp = processed[name]["TC"]["Air_out"]["p"] / processed[name]["TC"]["Air_in"]["p"]
        # Calculating the temperature after the compressor, based on ideal gas assumption
        processed[name]["TC"]["Air_out"]["T"] = processed[name]["TC"]["Air_in"]["T"] * (1 + beta_comp**((
            CONSTANTS["General"]["K_AIR"]-1)/CONSTANTS["General"]["K_AIR"])) / comp_isentropic_efficiency
        # The state of the air at the outlet of the compressor, outlet of the TC block, and inlet of the bypass valve
        #  are the same
        processed[name]["Comp"]["Air_out"]["T"] = processed[name]["TC"]["Air_out"]["T"]
        processed[name]["BPvalve"]["Air_in"]["T"] = processed[name]["TC"]["Air_out"]["T"]
        # The same temperature is the inlet to the charge air cooler
        processed[name]["CAC_HT"]["Air_in"]["T"] = processed[name]["TC"]["Air_out"]["T"]
        # Calculating the air inflow aspired by the cylinder: calculated as inlet air density times the maximum volume,
        # times the engine speed
        processed[name]["Cyl"]["Air_in"]["mdot"] = CONSTANTS["MainEngines"]["V_MAX"] * (
            processed[name]["Cyl"]["Air_in"]["p"] * 1e5) / (
            CONSTANTS["General"]["R_AIR"] / processed[name]["Cyl"]["Air_in"]["T"]) * (
            processed[name]["Cyl"]["Power_out"]["omega"] / 60 / 2 * CONSTANTS["General"]["ETA_VOL"]) * (
            CONSTANTS["AuxEngines"]["N_CYL"])
        processed[name]["Cyl"]["EG_out"]["mdot"] = processed[name]["Cyl"]["Air_in"]["mdot"] + processed[name]["Cyl"][
            "FuelPh_in"]["mdot"]
        # Here we calculate the bypass mass flow. THIS NEEDS TO BE CHECKED FOR CONSISTENCY (i.e. the bypass mass flow
        #  should always be kind of low, and anyway only be seen at low to medium load).
        # The equation is the result of a mass and energy balance over the whole engine
        processed[name]["BPvalve"]["Air_in"]["mdot"] = (
            processed[name]["Cyl"]["EG_out"]["mdot"] * CONSTANTS["General"]["CP_EG"] * (
                processed[name]["TC"]["EG_in"]["T"] - processed[name]["TC"]["EG_out"]["T"]) * CONSTANTS["MainEngines"]["ETA_MECH_TC"] -
            CONSTANTS["General"]["CP_AIR"] * processed[name]["Cyl"]["Air_in"]["mdot"] * (
                processed[name]["Comp"]["Air_out"]["T"] - processed[name]["Comp"]["Air_in"]["T"])) / (
            CONSTANTS["General"]["CP_AIR"] * (processed[name]["Comp"]["Air_out"]["T"] - processed[name]["Comp"]["Air_in"]["T"]) -
            CONSTANTS["General"]["CP_EG"] * CONSTANTS["MainEngines"]["ETA_MECH_TC"] * (
                processed[name]["TC"]["EG_in"]["T"] - processed[name]["TC"]["EG_out"]["T"]))
        processed[name]["BPvalve"]["Air_out"]["mdot"] = processed[name]["BPvalve"]["Air_in"]["mdot"]
        # The air mass flow going through the compressor is equal to the sum of the air flow through the bypass valve and
        # to the cylinders
        processed[name]["TC"]["Air_in"]["mdot"] = processed[name]["BPvalve"]["Air_in"]["mdot"] + processed[name]["Cyl"]["Air_in"]["mdot"]
        # Inlet and outlet flows to the compressor are equal
        processed[name]["TC"]["Air_out"]["mdot"] = processed[name]["TC"]["Air_in"]["mdot"]
        # The flow through the turbine is equal to the sum of the bypass flow and the exhaust coming from the cylinders
        processed[name]["TC"]["EG_in"]["mdot"] = processed[name]["BPvalve"]["Air_in"]["mdot"] + processed[name]["Cyl"]["EG_out"]["mdot"]
        # Calculating the temperature of the mixture after the merge between bypass and exhaust gas from the cylinders
        processed[name]["Cyl"]["EG_out"]["T"] = 298 + (
            processed[name]["TC"]["EG_in"]["mdot"] * CONSTANTS["General"]["CP_EG"] * (processed[name]["TC"]["EG_in"]["T"] - 298) -
            processed[name]["BPvalve"]["Air_in"]["mdot"] * CONSTANTS["General"]["CP_AIR"] * (processed[name]["TC"]["Air_out"]["T"] - 298)) / (
            processed[name]["Cyl"]["EG_out"]["mdot"] * CONSTANTS["General"]["CP_EG"])
    return processed

def readOtherExistingValues(raw, processed):
    # Other components
    processed["Other"]["SWC13"]["SeaWater"]["T_out"] = raw["SWC13_SeaWater_Tout"]  # CHECK IF IT IS IN OR OUT
    processed["Other"]["SWC24"]["SeaWater"]["T_out"] = raw["SWC24_SeaWater_Tout"]  # CHECK IF IT IS IN OR OUT
    processed["Other"]["SWC24"]["SeaWater"]["T_in"] = raw["SeaWater_T"]
    processed["Other"]["SWC24"]["SeaWater"]["T_in"] = raw["SeaWater_T"]
    return processed

def engineStatusCalculation(type, raw, processed, CONSTANTS, status, hd):
    for name in CONSTANTS["General"]["NAMES"][type]:
        status[name]["Load"] = processed[name]["Cyl"]["Power_out"]["Wdot"] / CONSTANTS[type]["MCR"]
        # We consider that the engines are on if the RPM of the turbocharger is higher than 5000 RPM
        #status[name]["OnOff"] = (status[name]["Load"] > 0.05) & (processed[name]["Cyl"]["Power_out"]["omega"] > CONSTANTS["MainEngines"]["RPM_DES"] * 0.1)
        status[name]["OnOff"] = raw[hd[name+"-TC__RPM_"]] > 5000
    return status

def engineCoolingSystemsCalculation(processed, CONSTANTS, status, engine_type):
    # This function calculates the different flows related to the cooling systems of the main engines.
    for name in CONSTANTS["General"]["NAMES"][engine_type]:
        # Calculating the total energy flow going to the cooling systems, based on the energy balance on the engine
        energy_2_cooling = (processed[name]["Cyl"]["FuelCh_in"]["Wdot"] -
            processed[name]["Cyl"]["Power_out"]["Wdot"] +
            CONSTANTS["General"]["CP_HFO"] * processed[name]["Cyl"]["FuelPh_in"]["mdot"] *
                (processed[name]["Cyl"]["FuelPh_in"]["T"] - processed["T_0"]) -
            CONSTANTS["General"]["CP_EG"] * processed[name]["TC"]["EG_out"]["mdot"] *
                (processed[name]["TC"]["EG_out"]["T"] - processed["T_0"]) +
            CONSTANTS["General"]["CP_AIR"] * processed[name]["TC"]["Air_in"]["mdot"] *
                (processed[name]["TC"]["Air_in"]["T"] - processed["T_0"]))
        # Calculating the energy going to the charge air cooler, based on the estimated temperatures on the air line
        energy_2_cac = CONSTANTS["General"]["CP_AIR"] * processed[name]["Cyl"]["Air_in"]["mdot"] * (processed[name]["TC"]["Air_out"]["T"] - processed[name]["Cyl"]["Air_in"]["T"])
        # Calculating the energy going to the HT cooling systems, based on interpolation from the project guide
        energy_2_ht_theoric = status[name]["Load"].apply(piecewisePolyvalHelperFunction,args=(CONSTANTS[engine_type]["POLY_LOAD_2_QDOT_HT"],)) * CONSTANTS[engine_type]["QDOT_HT_DES"]
        energy_2_lt_theoric = status[name]["Load"].apply(piecewisePolyvalHelperFunction,args=(CONSTANTS[engine_type]["POLY_LOAD_2_QDOT_LT"],)) * CONSTANTS[engine_type]["QDOT_LT_DES"]
        # The values calculated based on the project guide are reconciled based on the energy balance
        energy_2_ht = energy_2_cooling * energy_2_ht_theoric / (energy_2_ht_theoric + energy_2_lt_theoric)
        energy_2_lt = energy_2_cooling - energy_2_ht
        # The energy going to the CAC, HT stage is calculated assuming a 85% effectiveness of the heat exchanger
        energy_2_cac_ht = CONSTANTS[engine_type]["EPS_CAC_HTSTAGE"] * processed[name]["TC"]["Air_in"]["mdot"] * CONSTANTS["General"]["CP_AIR"] * (
            processed[name]["TC"]["Air_out"]["T"] - processed[name]["CAC_HT"]["HTWater_in"]["T"])
        # The energy going to the CAC, LT stage results as a consequence by thermal balance over the CAC
        energy_2_cac_lt = energy_2_cac - energy_2_cac_ht
        # The energy to the JWC results as a balance over the HT cooling systems
        energy_2_jwc = energy_2_ht - energy_2_cac_ht
        # The energy to the LOC results as a balance over the LT cooling systems
        energy_2_loc = energy_2_lt - energy_2_cac_lt
        # For all pumps, it is here assumed that the flow scales only with the speed of the engine (biiiiig assumption)
        processed[name]["LOC"]["LTWater_in"]["mdot"] = CONSTANTS[engine_type]["MFR_LT"] * processed[name]["Cyl"]["Power_out"]["omega"] / CONSTANTS[engine_type]["RPM_DES"]
        processed[name]["JWC"]["HTWater_in"]["mdot"] = CONSTANTS[engine_type]["MFR_HT"] * processed[name]["Cyl"]["Power_out"]["omega"] / CONSTANTS[engine_type]["RPM_DES"]
        # Asssigning values based on mass balances
        processed[name]["LOC"]["LTWater_out"]["mdot"] = processed[name]["LOC"]["LTWater_in"]["mdot"]
        processed[name]["CAC_LT"]["LTWater_out"]["mdot"] = processed[name]["LOC"]["LTWater_in"]["mdot"]
        processed[name]["CAC_LT"]["LTWater_out"]["mdot"] = processed[name]["LOC"]["LTWater_in"]["mdot"]
        processed[name]["JWC"]["HTWater_out"]["mdot"] = processed[name]["JWC"]["HTWater_in"]["mdot"]
        processed[name]["CAC_HT"]["HTWater_in"]["mdot"] = processed[name]["JWC"]["HTWater_in"]["mdot"]
        processed[name]["CAC_HT"]["HTWater_out"]["mdot"] = processed[name]["JWC"]["HTWater_in"]["mdot"]
        # Finally, the temperatures in the flows are calculated based on the calculated energy and mass flow values
        # For LT, first we have the CAC, then the LOC
        processed[name]["CAC_LT"]["LTWater_out"]["T"] = processed[name]["CAC_LT"]["LTWater_in"]["T"] + energy_2_cac_lt / processed[name]["CAC_LT"]["LTWater_out"]["mdot"] / CONSTANTS["General"]["CP_WATER"]
        processed[name]["LOC"]["LTWater_in"]["T"] = processed[name]["CAC_LT"]["LTWater_out"]["T"]
        processed[name]["LOC"]["LTWater_out"]["T"] = processed[name]["LOC"]["LTWater_in"]["T"] + energy_2_loc / processed[name]["LOC"]["LTWater_out"]["mdot"] / CONSTANTS["General"]["CP_WATER"]
        # For HT, first we have the JWC, then the CAC, HT
        processed[name]["JWC"]["HTWater_out"]["T"] = processed[name]["JWC"]["HTWater_in"]["T"] + energy_2_jwc / processed[name]["JWC"]["HTWater_out"]["mdot"] / CONSTANTS["General"]["CP_WATER"]
        processed[name]["CAC_HT"]["HTWater_in"]["T"] = processed[name]["JWC"]["HTWater_out"]["T"]
        processed[name]["CAC_HT"]["HTWater_out"]["T"] = processed[name]["CAC_HT"]["HTWater_in"]["T"] + energy_2_cac_ht / processed[name]["CAC_HT"]["HTWater_out"]["mdot"] / CONSTANTS["General"]["CP_WATER"]
        # For the LOC, we know the outlet (lower) temperature, we calculate the inlet temperature
        processed[name]["LOC"]["LubOil_out"]["mdot"][:] = CONSTANTS[engine_type]["MFR_LO"]
        processed[name]["LOC"]["LubOil_in"]["mdot"][:] = CONSTANTS[engine_type]["MFR_LO"]
        processed[name]["LOC"]["LubOil_in"]["T"] = processed[name]["LOC"]["LubOil_out"]["T"] + energy_2_loc / processed[name]["LOC"]["LubOil_out"]["mdot"] / CONSTANTS["General"]["CP_LO"]
    return processed

def bsfcISOCorrection(bsfc_ISO, charge_air_temp, charge_air_cooling_temp, fuel_temp, CONSTANTS):
    # This function calculates the "real" BSFC starting from the ISO corrected one and from measurements of
    # - Charge air temperature [K]
    # - Charge air coolant temperature [K]
    # - Fuel LHV [MJ/kg]
    # - Mechanical efficiency (often assumed at 0.8)

    # Assigning the value of the LHV depending on the fuel temperature
    LHV = pd.Series(0,index=charge_air_temp.index)
    LHV[fuel_temp < 70] = CONSTANTS["General"]["LHV_MDO"] # If T_fuel<70, it is Diesel
    LHV[fuel_temp >= 70] = CONSTANTS["General"]["LHV_HFO"] # If T_fuel>70, it is HFO
    # Converting existing data (expected in the form of dataSeries
    if isinstance(charge_air_temp,pd.Series):
        T_ca = charge_air_temp.values
    else:
        print("Error: Expecting a pandas data series as data type")
    if isinstance(charge_air_cooling_temp,pd.Series):
        T_lt = charge_air_cooling_temp.values
    else:
        print("Error: Expecting a pandas data series as data type")
    # Providing reference values for the variables
    k = (CONSTANTS["General"]["ISO"]["T_CA"] / T_ca)**1.2 * (CONSTANTS["General"]["ISO"]["T_LT"] / T_lt)
    alpha = k - 0.7 * (1 - k) * (1/CONSTANTS["General"]["ISO"]["ETA_MECH"] - 1)
    beta = k / alpha
    # Final calculation of the BSFC
    bsfc = bsfc_ISO * CONSTANTS["General"]["ISO"]["LHV"] / LHV * beta
    return (bsfc, LHV)

def polyvalHelperFunction(x,p):
    # The problem with applying "polyval" to data series is that the "x" is the second argument of the function
    # instead of being the first. So we use this function to invert the two, waiting to find a better way
    output = np.polyval(p,x)
    return output

def piecewisePolyvalHelperFunction(x,p):
    # The problem with applying "polyval" to data series is that the "x" is the second argument of the function
    # instead of being the first. So we use this function to invert the two, waiting to find a better way
    output = np.piecewise(x, [x < 0.5 , x >= 0.5], [np.polyval(p[1],x) , np.polyval(p[0],x)])
    return output


# paper:
# 	yolo1: https://arxiv.org/abs/1506.02640
# 	yolo2: https://arxiv.org/abs/1612.08242
# 	yolo3: https://pjreddie.com/media/files/papers/YOLOv3.pdf
# Author: Charles
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import sys
sys.path.append('../')
from utils.cfg_utils import *
from utils.standard_utils import *
from loss_module.region_loss import RegionLoss
from layers.yolo_layer import YoloLayer
from cfg import *

# pad:
# 	pre-Tensor <Right and Bottom> expand 1 dim.
# 	(the same as pre-final <Right and Bottom>)
# maxpool2d:
# 	stride: 1
# 	k_size: (2, 2)
# Input:
# 	(batch_size, in_channels, w, h)
# Output:
# 	the same as Input.
class MaxPoolStride1(nn.Module):
	def __init__(self):
		super(MaxPoolStride1, self).__init__()
	def forward(self, x):
		x_pad = F.pad(x, (0, 1, 0, 1), mode='replicate')
		x = F.max_pool2d(x_pad, 2, stride=1)
		return x

# Function:
# 	expand H, W using their own elements
# Example(stride=2):
# 	(0, 0, ., .):
# 		origin:
# 			1 2
# 			3 4
# 		changed:
# 			1 1 2 2
# 			1 1 2 2
# 			3 3 4 4
# 			3 3 4 4
# Input:
# 	(batch_size, in_channels, h, w)
# Output:
# 	(batch_size, in_channels, h*stride, w*stride)
class Upsample(nn.Module):
	def __init__(self, stride=2):
		super(Upsample, self).__init__()
		self.stride = stride
	def forward(self, x):
		assert(x.data.dim() == 4)
		B = x.data.size(0)
		C = x.data.size(1)
		H = x.data.size(2)
		W = x.data.size(3)
		x = x.view(B, C, H, 1, W, 1).expand(B, C, H, self.stride, W, self.stride).contiguous().view(B, C, H*self.stride, W*self.stride)
		return x

# Function:
# 	(B, C, H, W) -> (B, hs*ws*C, H//hs, W//ws)
# 	Get each C using stride that feels like skipping
# Input:
# 	(batch_size, in_channels, h, w)
# Output:
# 	(batch_size, ws*hs*C, H//hs, W//ws)
class Reorg(nn.Module):
	def __init__(self, stride=2):
		super(Reorg, self).__init__()
		self.stride = stride
	def forward(self, x):
		assert(x.data.dim() == 4)
		B = x.data.size(0)
		C = x.data.size(1)
		H = x.data.size(2)
		W = x.data.size(3)
		assert(H % self.stride == 0)
		assert(W % self.stride == 0)
		w_stride = self.stride
		h_stride = self.stride
		# (B, C, H, W) -> (B, C, H//hs, W//ws, hs, ws)
		x = x.view(B, C, H//h_stride, h_stride, W//w_stride, w_stride).transpose(3, 4).contiguous()
		# (B, C, H//hs, W//ws, hs, ws) -> (B, C, hs*ws, -1)
		x = x.view(B, C, (H//h_stride)*(W//w_stride), h_stride*w_stride).transpose(2, 3).contiguous()
		# (B, C, hs*ws, -1) -> (B, hs*ws, C, H//hs, W//ws)
		x = x.view(B, C, h_stride*w_stride, H//h_stride, W//w_stride).transpose(1, 2).contiguous()
		# (B, hs*ws, C, H//hs, W//ws) -> (B, hs*ws*C, H//hs, W//ws)
		x = x.view(B, h_stride*w_stride*C, H//h_stride, W//w_stride)
		return x

# Function:
# 	compute all elements' average each (H, W)
# Input:
# 	(N, C, H, W)
# Output:
# 	(N, C)
class GlobalAvgPool2d(nn.Module):
	def __init__(self):
		super(GlobalAvgPool2d, self).__init__()
	def forward(self, x):
		N = x.data.size(0)
		C = x.data.size(1)
		H = x.data.size(2)
		W = x.data.size(3)
		x = F.avg_pool2d(x, (H, W))
		x = x.view(N, C)
		return x

# Function:
# 	for route and shortcut
class EmptyModule(nn.Module):
	def __init__(self):
		super(EmptyModule, self).__init__()
	def forward(self, x):
		return x

# Darknet
# Support route shortcut and reorg
class Darknet(nn.Module):
	def __init__(self, cfgfile):
		super(Darknet, self).__init__()
		self.is_yolo3 = False
		self.blocks = parse_cfg(cfgfile)
		self.losses = []
		self.models = self.create_network(self.blocks)
		# Because of all yolo layers' anchors, num_anchors, anchor_step and num_classes is the same.
		self.loss = self.models[len(self.models)-1]
		self.width = int(self.blocks[0]['width'])
		self.height = int(self.blocks[0]['height'])
		if self.blocks[(len(self.blocks)-1)]['type'] == 'region' or 'yolo':
			self.anchors = self.loss.anchors
			self.num_anchors = self.loss.num_anchors
			self.anchor_step = self.loss.anchor_step
			self.num_classes = self.loss.num_classes
		self.header = torch.IntTensor([0, 0, 0, 0, 0])
		self.seen = 0
		# self.multiGPU = False
	def forward(self, x):
		ind = -2
		is_yolo3 = self.is_yolo3
		yolo_outs = []
		outputs = dict()
		# out_boxes = []
		for block in self.blocks:
			ind += 1
			if block['type'] == 'net':
				continue
			elif block['type'] in ['convolutional', 'maxpool', 'reorg', 'upsample', 'avgpool', 'softmax', 'connected']:
				x = self.models[ind](x)
				outputs[ind] = x
			elif block['type'] == 'route':
				layers = block['layers'].split(',')
				layers = [int(i) if int(i) > 0 else int(i)+ind for i in layers]
				if len(layers) == 1:
					x = outputs[layers[0]]
					outputs[ind] = x
				elif len(layers) == 2:
					x1 = outputs[layers[0]]
					x2 = outputs[layers[1]]
					x = torch.cat((x1, x2), 1)
					outputs[ind] = x
			elif block['type'] == 'shortcut':
				from_layer = int(block['from'])
				activation = block['activation']
				from_layer = from_layer if from_layer > 0 else from_layer + ind
				x1 = outputs[from_layer]
				x2 = outputs[ind-1]
				x  = x1 + x2
				if activation == 'leaky':
					x = F.leaky_relu(x, 0.1, inplace=True)
				elif activation == 'relu':
					x = F.relu(x, inplace=True)
				outputs[ind] = x
			elif block['type'] == 'region':
				continue
				# if self.training:
				# 	yolo_outs.append(x)
				# 	outputs[ind] = None
				# else:
				# 	boxes = self.models[ind](x)
				# 	out_boxes.append(boxes)
			elif block['type'] == 'yolo':
				# continue
				is_yolo3 = True
				yolo_outs.append(x)
				# if self.training:
				# 	yolo_outs.append(x)
				# 	outputs[ind] = None
				# else:
				# 	boxes = self.models[ind](x)
				# 	out_boxes.append(boxes)
			elif block['type'] == 'cost':
				continue
			else:
				print('[Error]:unkown type %s' % (block['type']))
		# if self.training:
		# 	# return self.loss
		# 	if is_yolo3:
		# 		# print('[INFO]:This yolov3 darknet-train...')
		# 		return yolo_outs
		# 	else:
		# 		return x
		# else:
		# 	if is_yolo3:
		# 		# print('[INFO]:This yolov3 darknet-test...')
		# 		return out_boxes
		# 	else:
		# 		return x
		# return x
		if is_yolo3:
			return yolo_outs
		else:
			return x
	# merge conv, bn, leaky, etc.
	def create_network(self, blocks):
		models = nn.ModuleList()
		prev_filters = 3
		out_filters =[]
		prev_stride = 1
		out_strides = []
		conv_id = 0
		# all conv stride = 1
		# So out_strides record sizes reduce in scale
		for block in blocks:
			if block['type'] == 'net':
				prev_filters = int(block['channels'])
				continue
			elif block['type'] == 'convolutional':
				conv_id = conv_id + 1
				batch_normalize = int(block['batch_normalize'])
				filters = int(block['filters'])
				kernel_size = int(block['size'])
				stride = int(block['stride'])
				is_pad = int(block['pad'])
				# Problem:
				# 	kernel_size when odd better?
				pad = (kernel_size-1)//2 if is_pad else 0
				activation = block['activation']
				model = nn.Sequential()
				if batch_normalize:
					model.add_module('conv{}'.format(conv_id), nn.Conv2d(prev_filters, filters, kernel_size, stride, pad, bias=False))
					model.add_module('bn{}'.format(conv_id), nn.BatchNorm2d(filters))
				else:
					model.add_module('conv{}'.format(conv_id), nn.Conv2d(prev_filters, filters, kernel_size, stride, pad))
				if activation == 'leaky':
					model.add_module('leaky{}'.format(conv_id), nn.LeakyReLU(0.1, inplace=True))
				elif activation == 'relu':
					model.add_module('relu{}'.format(conv_id), nn.ReLU(inplace=True))
				prev_filters = filters
				out_filters.append(prev_filters)
				prev_stride = stride * prev_stride
				out_strides.append(prev_stride)
				models.append(model)
			elif block['type'] == 'maxpool':
				pool_size = int(block['size'])
				stride = int(block['stride'])
				if stride > 1:
					model = nn.MaxPool2d(pool_size, stride)
				else:
					model = MaxPoolStride1()
				out_filters.append(prev_filters)
				prev_stride = stride * prev_stride
				out_strides.append(prev_stride)
				models.append(model)
			# at end in general
			elif block['type'] == 'avgpool':
				model = GlobalAvgPool2d()
				out_filters.append(prev_filters)
				models.append(model)
			elif block['type'] == 'softmax':
				model = nn.Softmax()
				out_strides.append(prev_stride)
				out_filters.append(prev_filters)
				models.append(model)
			# the losses are averaged over observations for each minibatch.
			elif block['type'] == 'cost':
				if block['_type'] == 'sse':
					model = nn.MSELoss(size_average=True)
				elif block['_type'] == 'L1':
					model = nn.L1Loss(size_average=True)
				elif block['_type'] == 'smooth':
					model = nn.SmoothL1Loss(size_average=True)
				out_filters.append(1)
				out_strides.append(prev_stride)
				models.append(model)
			elif block['type'] == 'reorg':
				stride = int(block['stride'])
				prev_filters = stride * stride * prev_filters
				out_filters.append(prev_filters)
				prev_stride = prev_stride * stride
				out_strides.append(prev_stride)
				models.append(Reorg(stride))
			elif block['type'] == 'upsample':
				stride = int(block['stride'])
				out_filters.append(prev_filters)
				prev_stride = prev_stride // stride
				out_strides.append(prev_stride)
				models.append(Upsample(stride))
			elif block['type'] == 'route':
				layers = block['layers'].split(',')
				ind = len(models)
				layers = [int(i) if int(i) > 0 else int(i)+ind for i in layers]
				if len(layers) == 1:
					prev_filters = out_filters[layers[0]]
					prev_stride = out_strides[layers[0]]
				elif len(layers) == 2:
					assert(layers[0] == ind - 1)
					prev_filters = out_filters[layers[0]] + out_filters[layers[1]]
					prev_stride = out_strides[layers[0]]
				out_filters.append(prev_filters)
				out_strides.append(prev_stride)
				models.append(EmptyModule())
			# other params used in forward function
			elif block['type'] == 'shortcut':
				ind = len(models)
				prev_filters = out_filters[ind-1]
				out_filters.append(prev_filters)
				prev_stride = out_strides[ind-1]
				out_strides.append(prev_stride)
				models.append(EmptyModule())
			elif block['type'] == 'connected':
				filters = int(block['output'])
				if block['activation'] == 'linear':
					model = nn.Linear(prev_filters, filters)
				elif block['activation'] == 'leaky':
					model = nn.Sequential(
								nn.Linear(prev_filters, filters),
								nn.LeakyReLU(0.1, inplace=True))
				elif block['activation'] == 'relu':
					model = nn.Sequential(
								nn.Linear(prev_filters, filters),
								nn.ReLU(inplace=True))
				prev_filters = filters
				out_filters.append(prev_filters)
				out_strides.append(prev_stride)
				models.append(model)
			# Note:
			# 	some values aren't assigned to class
			# 	(RegionLoss and YoloLayer())
			elif block['type'] == 'region':
				loss = RegionLoss()
				anchors = block['anchors'].split(',')
				loss.anchors = [float(i) for i in anchors]
				loss.num_classes = int(block['classes'])
				loss.num_anchors = int(block['num'])
				loss.anchor_step = len(loss.anchors) // loss.num_anchors
				loss.object_scale = float(block['object_scale'])
				loss.noobject_scale = float(block['noobject_scale'])
				loss.class_scale = float(block['class_scale'])
				loss.coord_scale = float(block['coord_scale'])
				loss.stride = prev_stride
				out_filters.append(prev_filters)
				out_strides.append(prev_stride)
				models.append(loss)
				self.losses.append(loss)
			elif block['type'] == 'yolo':
				self.is_yolo3 = True
				yolo_layer = YoloLayer()
				anchors = block['anchors'].split(',')
				anchor_mask = block['mask'].split(',')
				yolo_layer.anchor_mask = [int(i) for i in anchor_mask]
				yolo_layer.anchors = [float(i) for i in anchors]
				yolo_layer.num_classes = int(block['classes'])
				yolo_layer.num_anchors = int(block['num'])
				yolo_layer.anchor_step = len(yolo_layer.anchors) // yolo_layer.num_anchors
				yolo_layer.stride = prev_stride
				# yolo_layer.object_scale = float(block['object_scale'])
				# yolo_layer.noobject_scale = float(block['noobject_scale'])
				# yolo_layer.class_scale = float(block['class_scale'])
				# yolo_layer.coord_scale = float(block['coord_scale'])
				out_filters.append(prev_filters)
				out_strides.append(prev_stride)
				models.append(yolo_layer)
				self.losses.append(yolo_layer)
			else:
				print('[Error]:unkown type %s' % (block['type']))
		return models
	def print_network(self):
		print_cfg(self.blocks)
	def load_weights(self, weightfile):
		fp = open(weightfile, 'rb')
		header = np.fromfile(fp, count=5, dtype=np.int32)
		# header = np.fromfile(fp, count=4, dtype=np.int32)
		self.header = torch.from_numpy(header)
		self.seen = self.header[3]
		# print(self.header)
		# print(self.seen)
		buf = np.fromfile(fp, dtype=np.float32)
		# print(len(buf))
		fp.close()
		start = 0
		ind = -2
		for block in self.blocks:
			if start >= buf.size:
				break
			ind = ind + 1
			if block['type'] == 'net':
				continue
			elif block['type'] == 'convolutional':
				model = self.models[ind]
				batch_normalize = int(block['batch_normalize'])
				if batch_normalize:
					start = load_conv_bn(buf, start, model[0], model[1])
				else:
					start = load_conv(buf, start, model[0])
			elif block['type'] == 'connected':
				model = self.models[ind]
				if block['activation'] != 'linear':
					start = load_fc(buf, start, model[0])
				else:
					start = load_fc(buf, start, model)
			# Easier to develop
			elif block['type'] == 'maxpool':
				continue
			elif block['type'] == 'reorg':
				continue
			elif block['type'] == 'upsample':
				continue
			elif block['type'] == 'route':
				continue
			elif block['type'] == 'shortcut':
				continue
			elif block['type'] == 'region':
				continue
			elif block['type'] == 'yolo':
				continue
			elif block['type'] == 'avgpool':
				continue
			elif block['type'] == 'softmax':
				continue
			elif block['type'] == 'cost':
				continue
			else:
				print('[Error]:unkown type %s' % (block['type']))
	# cutoff:
	# 	whether delete end block or not
	def save_weights(self, outfile, cutoff=0):
		if cutoff <= 0:
			cutoff = len(self.blocks) - 1
		fp = open(outfile, 'wb')
		self.header[3] = self.seen
		header = self.header
		header.numpy().tofile(fp)
		ind = -1
		for blockId in range(1, cutoff+1):
			ind = ind + 1
			block = self.blocks[blockId]
			if block['type'] == 'convolutional':
				model = self.models[ind]
				batch_normalize = int(block['batch_normalize'])
				if batch_normalize:
					save_conv_bn(fp, model[0], model[1])
				else:
					save_conv(fp, model[0])
			# activation fun(like ReLU) params don't need be saved.
			elif block['type'] == 'connected':
				model = self.models[ind]
				if block['activation'] != 'linear':
					save_fc(fc, model[0])
				else:
					save_fc(fc, model)
			# Easier to develop
			elif block['type'] == 'maxpool':
				continue
			elif block['type'] == 'reorg':
				continue
			elif block['type'] == 'upsample':
				continue
			elif block['type'] == 'route':
				continue
			elif block['type'] == 'shortcut':
				continue
			elif block['type'] == 'region':
				continue
			elif block['type'] == 'yolo':
				continue
			elif block['type'] == 'avgpool':
				continue
			elif block['type'] == 'softmax':
				continue
			elif block['type'] == 'cost':
				continue
			else:
				print('[Error]:unkown type %s' % (block['type']))
		fp.close()

if __name__ == '__main__':
	# x = torch.Tensor(2, 2, 4, 4).fill_(1)
	# x = nn.Parameter(x)
	# print(x)
	# print(GlobalAvgPool2d()(x))
	d = Darknet('../cfg/me/darknet19_wfc_face.cfg')
	# d = Darknet('../cfg/me/darknet19_wfc_face.cfg')
	d.print_network()
	d.load_weights('../weights/darknet19_wfc.weights')
	# d.load_weights(r'E:\GraduationProject\experiment_record\detection\yolov3_CelebA\train\weights\000042.weights')
	# d.load_weights('../weights/000050.weights')