#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Feb  6 15:15:34 2026

@author: stefanoghirlandi
"""
#%%
import sys

print(sys.executable)
print(sys.version)

import numpy
import scipy
import pandas
import matplotlib

print("Packages loaded successfully")

#%%
import numpy as np
import scipy.stats
import pandas as pd
from pathlib import Path
import matplotlib.pyplot as plt

reduction_on = True  # set to True to apply titanium intensity reduction

lte_on = True  # set to True to apply lifetime extension

path = Path(__file__).parent / 'input for code.xlsx'

planes = pd.read_excel(path, sheet_name="Stock of planes in EU27")
planes = planes.rename(columns={
    "Stock (nº of planes)": "stock_planes",
    "Stock (titanium ton)": "stock_ti",
    "Stock_reduced (titanium ton)": "stock_ti_red"
}).set_index('Year')
planes.index = planes.index.astype(int)

start_year, end_year = 1960, 2040
full_years = np.arange(start_year, end_year + 1)
n_years = len(full_years)

#%%
planes_projection = pd.DataFrame(index=full_years)

# === Backcast: Stock before 2000 ===
target_stock_2000 = planes.loc[2000, 'stock_planes']

def logistic(t, K, r, t0):
    return K / (1 + np.exp(-r * (t - t0)))

K = target_stock_2000
r = 0.07
t0 = 1988

backcast_stock = logistic(full_years, K, r, t0)

# Scale exactly to 2000
idx_2000 = np.where(full_years == 2000)[0][0]
backcast_stock *= target_stock_2000 / backcast_stock[idx_2000]

planes_projection['stock'] = 0.0

planes_projection.loc[1960:1999, 'stock'] = backcast_stock[full_years < 2000]

planes_projection.loc[planes.index, 'stock'] = planes['stock_planes']


planes_projection['inflow'] = 0.0
planes_projection['outflow'] = 0.0
planes_projection['nas'] = 0.0

switch_year = 2025  # year from which new inflows get extended lifetime

# Baseline survival curve
curve_old = scipy.stats.norm.sf(
    np.arange(n_years), loc=25, scale=12.5
)
curve_old /= curve_old[0]

# Extended lifetime survival curve (active only if lte_on)
curve_new = scipy.stats.norm.sf(
    np.arange(n_years), loc=30, scale=15
) if lte_on else curve_old
curve_new /= curve_new[0]

survival_matrix = np.zeros((n_years, n_years))

for i, cohort_year in enumerate(full_years):

    # Select survival curve based on inflow year
    if cohort_year < switch_year:
        curve = curve_old
    else:
        curve = curve_new

    survival_matrix[i:, i] = curve[:n_years - i]

survival_matrix_df = pd.DataFrame(
    survival_matrix, index=full_years, columns=full_years
)

# Compute inflows using stock-driven logic
for t, year in enumerate(full_years):

    if t == 0:
        planes_projection.at[year, 'inflow'] = planes_projection.loc[year, 'stock']
        continue

    survivors = survival_matrix_df.loc[year, full_years[:t]].dot(
        planes_projection.loc[full_years[:t], 'inflow']
    )
    stock_required = planes_projection.loc[year, 'stock']
    inflow = max(stock_required - survivors, 0)
    planes_projection.at[year, 'inflow'] = inflow

# Compute NAS and Outflow using continuity identity
planes_projection['nas'] = planes_projection['stock'].diff().fillna(0)
planes_projection['prev_stock'] = planes_projection['stock'].shift(1, fill_value=0)
planes_projection['outflow'] = (
    planes_projection['prev_stock'] + planes_projection['inflow'] - planes_projection['stock']
).clip(lower=0)
planes_projection.drop(columns=['prev_stock'], inplace=True)

planes_projection[['stock']].plot()
planes_projection[['inflow','outflow']].plot()

# Plane stock by vintage (cohort survival matrix in number of planes)

plane_stock_by_vintage = survival_matrix_df.multiply(
    planes_projection['inflow'],
    axis=1
)

#%% Titanium survival matrix

# --- Use reduced titanium intensity stock if reduction is on ---
stock_col = 'stock_ti_red' if reduction_on else 'stock_ti'

target_ti_2000 = planes.loc[2000, stock_col]

ti_projection = pd.DataFrame(index=full_years)
ti_projection['stock'] = 0.0
ti_projection['inflow'] = 0.0
ti_projection['outflow'] = 0.0
ti_projection['nas'] = 0.0

K = target_ti_2000
r = 0.07
t0 = 1988
backcast_ti = logistic(full_years, K, r, t0)
idx_2000 = np.where(full_years == 2000)[0][0]
backcast_ti *= target_ti_2000 / backcast_ti[idx_2000]

ti_projection.loc[1960:1999, 'stock'] = backcast_ti[full_years < 2000]
ti_projection.loc[planes.index, 'stock'] = planes[stock_col]

for t, year in enumerate(full_years):
    if t == 0:
        ti_projection.at[year, 'inflow'] = ti_projection.loc[year, 'stock']
        continue
    survivors = survival_matrix_df.loc[year, full_years[:t]].dot(
        ti_projection.loc[full_years[:t], 'inflow']
    )
    stock_required = ti_projection.loc[year, 'stock']
    inflow = max(stock_required - survivors, 0)
    ti_projection.at[year, 'inflow'] = inflow

ti_projection['nas'] = ti_projection['stock'].diff().fillna(0)
ti_projection['prev_stock'] = ti_projection['stock'].shift(1, fill_value=0)
ti_projection['outflow'] = (
    ti_projection['prev_stock']
    + ti_projection['inflow']
    - ti_projection['stock']
).clip(lower=0)
ti_projection.drop(columns=['prev_stock'], inplace=True)

ti_stock_by_vintage = survival_matrix_df.multiply(
    ti_projection['inflow'],
    axis=1
)

ti_per_plane = ti_projection['stock'] / planes_projection['stock']
ti_per_plane.plot()
#%% Recycling scenario
# Alpha is fixed at baseline value (23.3) throughout.
# Progressive import substitution is fully captured by the explicit shares.
# Reprocessing share rises from baseline 5% to scenario 20% (constant from 2025)
# End-of-life recycling share ramps linearly from 0% (2025) to 6% (2040)
# Annual displacement = (additional reprocessing + recycling) x baseline inflow x alpha_baseline

# --- Parameters ---
alpha_baseline        = 23.3   # upstream import coefficient, fixed at 2019 baseline (kt/kt)
reprocess_share       = 0.20   # scenario reprocessing share of imports, constant from 2025
reprocess_base_share  = 0.05   # baseline reprocessing share (2019 static MFA)
recycle_start_share   = 0.00   # end-of-life recycling share 2025
recycle_end_share     = 0.05   # end-of-life recycling share 2040

scenario_years = np.arange(2025, 2041)
n = len(scenario_years)

# --- Shares ---
rep_shares = np.full(n, reprocess_share)                                    # constant at 20%
rec_shares = np.linspace(recycle_start_share, recycle_end_share, n)        # linear ramp

recycling = pd.DataFrame(index=scenario_years)
recycling['inflow']  = ti_projection.loc[scenario_years, 'inflow'].values
recycling['outflow'] = ti_projection.loc[scenario_years, 'outflow'].values
recycling['imports'] = recycling['inflow'] * alpha_baseline                 # fixed coefficient

recycling['reprocessed']          = recycling['imports'] * rep_shares
recycling['reprocessed_baseline'] = recycling['imports'] * reprocess_base_share
recycling['recycled']             = recycling['imports'] * rec_shares

# --- Storage: cumulative retired aircraft material available for recycling ---
storage_init = ti_projection.loc[2000:2024, 'outflow'].sum() * 0.8
recycling['storage'] = 0.0
s = storage_init
for year in scenario_years:
    out = ti_projection.loc[year, 'outflow']
    eol = recycling.loc[year, 'recycled']
    s += out * 0.8       
    s -= eol / 0.8       
    recycling.loc[year, 'storage'] = s

# --- Annual displacement: (Δφ_rep + φ_rec,t) x I_baseline,t x alpha_baseline ---
recycling['annual_displacement'] = (
    (recycling['reprocessed'] - recycling['reprocessed_baseline'])
    + recycling['recycled']
)
recycling['cumulative_displacement'] = recycling['annual_displacement'].cumsum()

# --- Summary ---
print(recycling[['imports', 'reprocessed', 'recycled', 'storage',
                  'annual_displacement', 'cumulative_displacement']].round(1))

# Sanity check: storage should never go negative
if (recycling['storage'] < 0).any():
    first_negative = recycling[recycling['storage'] < 0].index[0]
    print(f"\nWARNING: storage depleted in {first_negative} — reduce recycle_end_share")
else:
    print(f"\nStorage remains positive throughout. Minimum: {recycling['storage'].min():.1f} t")

print(f"\nCumulative import displacement by 2040: "
      f"{recycling.loc[2040, 'cumulative_displacement']:.1f} t "
      f"({recycling.loc[2040, 'cumulative_displacement'] / 36_500 * 100:.1f}% of 2019 imports)")

#%%# --- Plotting ---

this_folder = Path(__file__).parent/'Plots'

# 1. Inflow, Outflow, NAS for planes
planes_projection.loc[2000:, ['inflow', 'outflow', 'nas']].plot(figsize=(12, 6))
plt.title('Aircraft Inflow, Outflow, and NAS')
plt.xlabel('Year')
plt.ylabel('Number of Planes')
plt.tight_layout()
plt.savefig(this_folder / '1. LTE - planes_inflow_outflow_nas.png', dpi=300)
plt.close()

# 2. Stock for planes
planes_projection.loc[2000:, 'stock'].plot(figsize=(12, 6))

plt.title('Aircraft Stock Over Time')
plt.xlabel('Year')
plt.ylabel('Number of Planes')
plt.tight_layout()
#plt.show()
plt.savefig(this_folder / '2. LTE - planes_stock_over_time.png', dpi=300)
plt.close()

# 3. Survival Matrix Area Plot for planes
plane_stock_by_vintage.loc[2000:, 2000:].plot(
    kind="area",
    stacked=True,
    legend=False,
    figsize=(12, 6)
)

plt.title('Aircraft Survival Matrix')
plt.xlabel('Year')
plt.ylabel('Number of planes')
plt.tight_layout()
plt.savefig(this_folder / '3. LTE - aircraft_survival_matrix.png', dpi=600)
plt.close()

# 4. Survival Matrix Area Plot for Titanium
ti_stock_by_vintage.loc[2000:, 2000:].plot(kind="area", stacked=True, legend=False,  figsize=(12, 6))
plt.title('Titanium Survival Matrix')
plt.xlabel('Year')
plt.ylabel('Titanium (tons)')
plt.tight_layout()
#plt.show()
plt.savefig(this_folder / '4. LTE - titanium_survival_matrix.png', dpi=300)
plt.close()

# --- Titanium inflow and outflow line plot ---
plt.figure(figsize=(10, 5))
ti_projection.loc[2000:, ['inflow', 'outflow']].plot(ax=plt.gca())
plt.title('Titanium Inflow and Outflow Over Time', fontsize=14)
plt.xlabel('Year')
plt.ylabel('Titanium (tons/year)')
plt.legend(['Inflow', 'Outflow'])
plt.grid(False)
plt.tight_layout()
plt.savefig(this_folder / '5. LTE - titanium_inflow_outflow.png', dpi=300)
plt.close()


plt.figure(figsize=(10, 5))
ti_projection.loc[2000:,'stock'].plot()
plt.title('Total Titanium Stock in Aircraft Over Time', fontsize=14)
plt.xlabel('Year')
plt.ylabel('Titanium in stock (tons)')
plt.grid(False)
plt.tight_layout()
plt.savefig(this_folder / '6. LTE - titanium_total_stock.png', dpi=300)
plt.close()

# Create figure with two subplots
fig, axes = plt.subplots(2, 1, figsize=(12, 10), sharex=False)

# --- (2) Inflow and Outflow Line Plot ---
ti_projection.loc[2000:,['inflow', 'outflow']].plot(
    ax=axes[0]
)
axes[0].set_title('Titanium inflow and outflow in the EU aircraft fleet (2000–2040)', fontsize=14)
axes[0].set_xlabel('Year')
axes[0].set_ylabel('Titanium (tons/year)')
axes[0].legend(['Inflow', 'Outflow'])
axes[0].grid(False)

# --- (1) Survival Matrix Area Plot ---
ti_stock_by_vintage.loc[2000:, 2000:].plot(
    kind="area",
    stacked=True,
    legend=False,
    ax=axes[1]
)
axes[1].set_title('Titanium stock evolution by vintage within the EU aircraft fleet (2000–2040)', fontsize=14)
axes[1].set_xlabel('Year')
axes[1].set_ylabel('Titanium (tons)')
axes[1].grid(False)

plt.tight_layout()
plt.savefig(this_folder / 'LTE - Combined_Titanium_Plots.png', dpi=300)
plt.close()

#%% Evolution of the fleet by classes of Ti intensity

# --- Import ---
shares = pd.read_excel(path, sheet_name="share of models per year", index_col=0)  # years x models
ti_content = pd.read_excel(path, sheet_name="Tonnage per model", index_col=0)  # models x 1

shares.index = shares.index.astype(int)

# Align models across both dataframes
common_models = shares.columns.intersection(ti_content.index)
shares = shares[common_models]
ti_content = ti_content.loc[common_models, 'Titanium (ton)']

# --- Classify aircraft by titanium content ---
bins = [0, 5, 10, 15, float('inf')]
labels = ['< 5 t', '5–10 t', '10–15 t', '> 15 t']

ti_class = pd.cut(ti_content, bins=bins, labels=labels, right=False)

# --- Sum shares by class for each year ---
class_shares = pd.DataFrame(index=shares.index)
for cls in labels:
    models_in_class = ti_class[ti_class == cls].index
    class_shares[cls] = shares[models_in_class].sum(axis=1)

# --- Plot ---
fig, ax = plt.subplots(figsize=(12, 6))

colors = ['#4c9be8', '#f4a442', '#d9534f', '#5cb85c']
for cls, color in zip(labels, colors):
    ax.plot(class_shares.index, class_shares[cls], label=cls, color=color, linewidth=1.8)

#ax.axvspan(2020, 2022, alpha=0.12, color='grey', label='COVID shock')
ax.axvline(2024, color='grey', linestyle='--', linewidth=0.8, label='Forecast boundary')

ax.set_title('Share of EU fleet titanium stock by aircraft titanium class (2000–2040)')
ax.set_xlabel('Year')
ax.set_ylabel('Share of total fleet titanium stock')
ax.legend(title='Ti content per aircraft', loc='upper center', bbox_to_anchor=(0.5, -0.15), ncols=6)
ax.set_xlim(shares.index.min(), shares.index.max())

plt.tight_layout()
plt.savefig(this_folder /'ti_class_shares.png', dpi=150)
plt.show()

#%% Save all planes and titanium yearly stock-flow in Excel

results_folder = Path(__file__).parent / 'Results'
results_folder.mkdir(exist_ok=True)

# --- Rounding helpers ---
def round_int(df):
    return df.round(0).astype('Int64')

# Aircraft counts (must be integers)
rounded_planes_projection = round_int(planes_projection)
rounded_plane_stock_by_vintage = round_int(plane_stock_by_vintage)
# Titanium quantities (tons): round for reporting ONLY
rounded_ti_projection = ti_projection.round(0)
rounded_ti_stock_by_vintage = ti_stock_by_vintage.round(0)

# --- Excel export ---
with pd.ExcelWriter(
    results_folder / 'Reduction+LTE - Aircraft fleet and titanium content.xlsx',
    engine='xlsxwriter'
) as writer:

    rounded_planes_projection.to_excel(writer, sheet_name='planes_projection')
    rounded_plane_stock_by_vintage.to_excel(writer, sheet_name='plane_stock_by_vintage')
    rounded_ti_projection.to_excel(writer, sheet_name='titanium_projection')
    rounded_ti_stock_by_vintage.to_excel(writer, sheet_name='titanium_stock_by_vintage')

    workbook = writer.book

    for sheet_name in writer.sheets:
        worksheet = writer.sheets[sheet_name]
        worksheet.freeze_panes(1, 1)
        worksheet.set_zoom(90)
        worksheet.set_column('A:ZZ', 15)