from pprint import pp
import logging
import configparser
import pathlib
import subprocess
from time import time

import pandas as pd
import numpy as np
from scipy.interpolate import griddata

_t0 = time()
config = configparser.ConfigParser()
config.read('config.ini')

if config['main']['plot'] == 'True': import matplotlib.pyplot as plt

logging.basicConfig(level=config['main']['logging'])
log = logging.getLogger('bathymetry')


# --- Initialization
log.info('Initialization')
folders = {
    'data': pathlib.Path(config['data']['root']),
    'scad': pathlib.Path(config['scad']['root']),
    'stl': pathlib.Path(config['stl']['root']),
    'pandas': pathlib.Path(config['pandas']['root']),
}
if config.getboolean('fig', 'enable', fallback=False):
    folders['fig'] = pathlib.Path(config.get('fig', 'root'))
for path in folders.values():
    path.mkdir(exist_ok=True)

bathy = pd.read_csv(
    folders['data'].joinpath(config['data']['bathy']),
    sep=' ',
    names=('lon', 'lat', 'z'),
)
artha = pd.read_csv(
    folders['data'].joinpath(config['data']['artha']),
)

earth_radius = 6371e3
L0 = float(config['main']['L0'])
L1 = float(config['main']['L1'])
direction = float(config['main']['dir'])

artha = pd.concat((artha, pd.DataFrame([
    {
        'lat': artha.lat.at[0] \
            + np.cos(direction*np.pi/180) * L0/earth_radius * 180/np.pi,
        'lon': artha.lon.at[0] \
            + np.sin(direction*np.pi/180) * L0/earth_radius * 180/np.pi \
                / np.cos(artha.lat.at[0] * np.pi/180),
    },
    {
        'lat': artha.lat.at[0] \
            - np.cos(direction*np.pi/180) * L1/earth_radius * 180/np.pi,
        'lon': artha.lon.at[0] \
            - np.sin(direction*np.pi/180) * L1/earth_radius * 180/np.pi \
                / np.cos(artha.lat.at[0] * np.pi/180),
    },
], index=[-1, 1])))

# --- Interpolation
log.info('Interpolating data')
N = int(config['main']['N'])
line = pd.DataFrame({
    'x': np.linspace(0, L0+L1, N),
    'lat': np.linspace(artha.lat.at[-1], artha.lat.at[1], N),
    'lon': np.linspace(artha.lon.at[-1], artha.lon.at[1], N),
})
line.set_index('x', inplace=True)

line['z'] = griddata(
    bathy[['lon','lat']],
    bathy.z,
    line[['lon','lat']],
    method='linear',
)

# --- Adding blocs
log.info('Adding blocs')
b = config['data']['blocs'].split(',')
blocs = pd.Series([
    pd.read_csv(
        folders['data'].joinpath(path),
        index_col='x',
    ) for path in b
])
for bloc in blocs:
    bloc.index = L0 - bloc.index


#line.z = all_blocs.z.fillna(line.z)

lim = float(config['main']['bathy_max'])
bathy_line = line.z.copy()
#bathy_line = line.z.clip(upper=lim)
bathy_line = bathy_line[(bathy_line.index - L0 < -10) | (bathy_line.index - L0 > 0)]
bathy_line.index = np.where((bathy_line > lim) & (bathy_line.index < L0), bathy_line.index + 10, bathy_line.index)
rubble_line = line.z[(line.z > lim) & (line.index < L0)]

data_dict = {}
for i in range(blocs.size):
    data_dict[f'bloc{i}'] = blocs.iat[i].reset_index().values
for name, data in (
        ('bathy', bathy_line),
        ('rubble', rubble_line),
    ):
    data_dict[name] = np.concatenate((
        data.reset_index().values,
        [[data.index.max(), data.min()],
        [data.index.min(),data.min()]]
    ))

# --- Generating SCAD

with open(folders['data'].joinpath(config['data']['base_scad'])) as bsf:
    base_scad = bsf.read()
for (name, data) in data_dict.items():
    log.info(f'Generating {name}')
    log.info('\tGenerating SCAD file')
    scad_file = folders['scad'].joinpath(f'{name}.scad')
    with open(scad_file, 'w') as osf:
        osf.write(base_scad.format(
            np.array2string(
                data,
                threshold=np.inf,
                separator=','
            )
        ))
    log.info('\tGenerating STL file')
    subprocess.run(('openscad', scad_file, '-o',
        folders['stl'].joinpath(f'{name}.stl')),
        check=True,
        capture_output=True)


# --- Saving pandas
log.info('Saving Pandas')

with pd.HDFStore(
        folders['pandas'].joinpath(config['pandas']['file']),
        mode='w',
        complib='blosc',
    ) as hdf:
    for name, data in (
            ('bathy', bathy_line),
            ('rubble', rubble_line),
        ):
        hdf.put(name, data)
    for i, bloc in blocs.items():
        hdf.put(f'bloc{i}', bloc)


_t1 = time()
log.info(f'Program ended successfully after {_t1-_t0:.2f}s')

# --- Plotting

if config['main']['plot'] == 'True':
    log.info('Plotting data')
    flt = (
        ((bathy.lon-artha.lon.at[0]).abs() < 0.002) & \
        ((bathy.lat-artha.lat.at[0]).abs() < 0.002)
    )
    
    fig, ax = plt.subplots()
    c = ax.tricontourf(
        bathy.lon[flt],
        bathy.lat[flt],
        bathy.z[flt],
        #marker='1',
        #lw=1,
    )
    fig.colorbar(c, label='Altitude (m)')
    ax.plot(
        artha.lon,
        artha.lat,
        color='k',
        marker='+',
        lw=1,
        label='Studied domain',
    )
    ax.legend()
    ax.grid(alpha=.2)
    ax.set(
        xlabel='Longitude (°)',
        ylabel='Latitude (°)',
        aspect='equal',
    )
    if config.getboolean('fig', 'enable', fallback=False):
        fig.savefig(folders['fig'].joinpath(
            f'map.{config.get("fig", "format")}'))
    
    fig, ax = plt.subplots(figsize=(6,2))
    ax.plot(
        bathy_line.index,
        bathy_line,
        color='k',
        lw=1,
        zorder=10,
        label='Bathymetry',
    )
    ax.plot(
        rubble_line.index,
        rubble_line,
        color='r',
        lw=1,
        zorder=11,
        label='Armour',
    )
    blocs.apply(lambda bloc: ax.fill_between(
        bloc.index,
        bloc.z,
        color='k',
        zorder=12,
        #alpha=.1,
    ))
    ax.set(
        xlabel='x (m)',
        ylabel='y (m)',
        aspect='equal',
        xlim=(bathy_line.index.min(), bathy_line.index.max()),
    )
    ax.grid()
    ax.legend()
    if config.getboolean('fig', 'enable', fallback=False):
        fig.savefig(folders['fig'].joinpath(
            f'bathy.{config.get("fig", "format")}'))
    
    plt.show(block=True)