# -*- coding: utf-8 -*-
"""
Created on Wed Dec  6 11:19:43 2023

@author: nielbock
"""
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy.interpolate import InterpolatedUnivariateSpline
import matplotlib.ticker as ticker
import os

# Get the directory
wdir=os.getcwd()

LANG="DE"
#LANG="EN"

EARTHMASS_SI = 5.97217e24
EARTHRADIUS_SI = 6371e3
DOTS = 100000

def sphere_radius(mass, bulk_density):
    mass_E = mass / EARTHMASS_SI
    volume = mass / bulk_density
    radius = (3 * volume / (4 * np.pi)) ** (1 / 3)
    radius_E = radius / EARTHRADIUS_SI
    return mass_E,radius_E

# Mass range for the plot in kg
mass = np.linspace(1e22, 1e27, DOTS)  # Adjust the range as needed

# Bulk densities for different functions in kg/m^3
bulk_densities = [996.557, 4107, 5513, 7874]  # Add or modify densities as desired

# Custom labels for substances
if LANG == "DE":
    labels = ['Wasser (konstantes Volumen)', 'Silikate (konstantes Volumen)', 'Erde (konstantes Volumen)', 'Eisen (konstantes Volumen)']
else:
    labels = ['Water (incompressible)', 'Silicate (incompressible)', 'Earth (incompressible)', 'Iron (incompressible)']

# Custom colors for curves
colors = ['blue', 'grey', 'green','red']

# Create a figure and axes
fig, ax = plt.subplots()

# Plot the radius for each bulk density
for density, label, color in zip(bulk_densities, labels, colors):
    mass_E,radius_E = sphere_radius(mass, density)
    ax.plot(mass_E, radius_E, color=color)



# Read empirical data from CSV file
data = pd.read_csv('massradius_100percentH2O_1000K_1mbar.txt')

# Extract mass and radius from empirical data (Earth values)
empirical_mass = data['mass']
empirical_radius = data['radius']

# Perform spline interpolation on empirical data
spline = InterpolatedUnivariateSpline(empirical_mass, empirical_radius)

# Generate interpolated dataset
interpolated_mass = np.linspace(empirical_mass.min(), empirical_mass.max(), DOTS)  # Adjust the number of points as needed
interpolated_radius = spline(interpolated_mass)

# Plot interpolated data
if LANG == "DE":
    model="Wasser bei 1000 K"
else:
    model="Water @ 1000 K"
ax.plot(interpolated_mass, interpolated_radius, color='cyan', linestyle='dashed', label=model)


# Read empirical data from CSV file
data = pd.read_csv('massradius_100percentH2O_300K_1mbar.txt')

# Extract mass and radius from empirical data (Earth values)
empirical_mass = data['mass']
empirical_radius = data['radius']

# Perform spline interpolation on empirical data
spline = InterpolatedUnivariateSpline(empirical_mass, empirical_radius)

# Generate interpolated dataset
interpolated_mass = np.linspace(empirical_mass.min(), empirical_mass.max(), DOTS)  # Adjust the number of points as needed
interpolated_radius = spline(interpolated_mass)

# Plot interpolated data
if LANG == "DE":
    model="Wasser bei 300 K"
else:
    model="Water @ 300 K"
ax.plot(interpolated_mass, interpolated_radius, color='blue', linestyle='dashed', label=model)


# Read empirical data from CSV file
data = pd.read_csv('massradiusmgsio3.txt')

# Extract mass and radius from empirical data (Earth values)
empirical_mass = data['mass']
empirical_radius = data['radius']

# Perform spline interpolation on empirical data
spline = InterpolatedUnivariateSpline(empirical_mass, empirical_radius)

# Generate interpolated dataset
interpolated_mass = np.linspace(empirical_mass.min(), empirical_mass.max(), DOTS)  # Adjust the number of points as needed
interpolated_radius = spline(interpolated_mass)

# Plot interpolated data
if LANG == "DE":
    model='Silikatplanet (MgSiO$_3$)'
else:
    model='Silicate Planet (MgSiO$_3$)'
ax.plot(interpolated_mass, interpolated_radius, color='grey', linestyle='dashed', label=model)


# Read empirical data from CSV file
data = pd.read_csv('massradiusEarthlikeRocky.txt')

# Extract mass and radius from empirical data (Earth values)
empirical_mass = data['mass']
empirical_radius = data['radius']

# Perform spline interpolation on empirical data
spline = InterpolatedUnivariateSpline(empirical_mass, empirical_radius)

# Generate interpolated dataset
interpolated_mass = np.linspace(empirical_mass.min(), empirical_mass.max(), DOTS)  # Adjust the number of points as needed
interpolated_radius = spline(interpolated_mass)

# Plot interpolated data
if LANG == "DE":
    model='Erdähnlicher Felsenplanet'
else:
    model='Earth-like Rocky Planet'
ax.plot(interpolated_mass, interpolated_radius, color='green', linestyle='dashed', label=model)


# Read empirical data from CSV file
data = pd.read_csv('massradiusFe.txt')

# Extract mass and radius from empirical data (Earth values)
empirical_mass = data['mass']
empirical_radius = data['radius']

# Perform spline interpolation on empirical data
spline = InterpolatedUnivariateSpline(empirical_mass, empirical_radius)

# Generate interpolated dataset
interpolated_mass = np.linspace(empirical_mass.min(), empirical_mass.max(), DOTS)  # Adjust the number of points as needed
interpolated_radius = spline(interpolated_mass)

# Plot interpolated data
if LANG == "DE":
    model='Eisenplanet'
else:
    model='Iron Planet'
ax.plot(interpolated_mass, interpolated_radius, color='red', linestyle='dashed', label=model)


if LANG == "DE":
    text = "Modelle: Zeng et al. (2019), DOI: 10.1073/pnas.181295116, https://lweb.cfa.harvard.edu/~lzeng/planetmodels.html"
else:    
    text = "Models: Zeng et al. (2019), DOI: 10.1073/pnas.181295116, https://lweb.cfa.harvard.edu/~lzeng/planetmodels.html"

ax.text(0.98, 0.02, text, transform=ax.transAxes, ha='right', va='bottom',fontsize=0.5 * plt.rcParams['font.size'])


# Plot Formula Mass vs. Radius for constant density
formula = r'$M = \frac{4}{3}\,\pi\,R^3\,\rho(R)$'
fig.text(0.15, 0.15, formula, transform=ax.transAxes, ha='left', va='bottom',fontsize=1.0 * plt.rcParams['font.size'])


# Set labels and title
if LANG == "DE":
    ax.set_xlabel('Masse (Erdmassen)')
    ax.set_ylabel('Radius (Erdradien)')
    ax.set_title('Planetenradius vs. Masse')
else:
    ax.set_xlabel('Mass (Earth mass)')
    ax.set_ylabel('Radius (Earth radius)')
    ax.set_title('Planet Radius vs. Mass')

# Set custom axis limits
x_min, x_max = 0, 10  # Custom x-axis limits
y_min, y_max = 0, 3.5  # Custom y-axis limits
ax.set_xlim(x_min, x_max)
ax.set_ylim(y_min, y_max)
ax.set_xticks(np.linspace(0, 10, 11))
ax.set_yticks(np.arange(0, 3.6, 0.5))

# Change x-axis tick format from point to comma for decimals
if LANG == "DE":
    # formatter = ticker.FuncFormatter(lambda x, pos: f'{x:.1f}'.replace('.', ','))
    # plt.gca().xaxis.set_major_formatter(formatter)
    formatter = ticker.FuncFormatter(lambda y, pos: f'{y:.1f}'.replace('.', ','))
    plt.gca().yaxis.set_major_formatter(formatter)

# Add legend with custom labels
ax.legend()

# Add legend with custom labels
legend_x = 0.98
legend_y = 0.05
ax.legend(loc='lower right', bbox_to_anchor=(legend_x, legend_y), fontsize=0.7 * plt.rcParams['font.size'])

# Export the plot to PNG and PDF with high resolution
dpi = 300  # Resolution in dots per inch

output_path = rf'{wdir}\Plots\Exoplanets_M_vs_R_combined'

if LANG == "DE":
    plt.savefig(f'{output_path}_de.png', dpi=dpi)
    plt.savefig(f'{output_path}_de.pdf', dpi=dpi)
else:
    plt.savefig(f'{output_path}_en.png', dpi=dpi)
    plt.savefig(f'{output_path}_en.pdf', dpi=dpi)

# Display the plot
plt.show()