#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright 2018-2024 Institute of Computer Science of the Czech Academy of
# Sciences, Prague, Czech Republic. Authors: Pavel Krc, Martin Bures, Jaroslav
# Resler.
#
# This file is part of PALM-METEO.
#
# PALM-METEO is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# PALM-METEO is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# PALM-METEO. If not, see <https://www.gnu.org/licenses/>.

import numpy as np
import netCDF4

from palmmeteo.plugins import WritePluginMixin
from palmmeteo.logging import die, warn, log, verbose
from palmmeteo.config import cfg, ConfigError
from palmmeteo.runtime import rt
from palmmeteo.library import PalmPhysics

ax_ = np.newaxis

# Convert LOD2 to LOD1 if requested
lod2_to_lod = lambda lod, arr: arr.mean(axis=(-2,-1)) if lod == 1 else arr

class WritePlugin(WritePluginMixin):
    def _setup(self):
        match cfg.simulation.approximation:
            case 'anelastic':
                self.is_anelastic = True
            case 'boussinesq':
                self.is_anelastic = False
            case _:
                ConfigError('Unknown approximation', cfg.simulation, 'approximation')

    def check_config(self, *args, **kwargs):
        self._setup()

    def write_data(self, fout, *args, **kwargs):
        log('Writing data to dynamic driver')

        self._setup()

        dtdefault = cfg.output.default_precision
        filldefault = cfg.output.fill_value

        # shorthands
        fov = fout.variables
        def mkvar(vname, dims, lod=None, units=None, dtype=dtdefault,
                fill=filldefault, attrs_from=None, orig_lod=None):
            if orig_lod==2 and lod==1:
                # Artificially remove y,x from dimensions
                dims = dims[:-2]

            v = fout.createVariable(vname, dtype, dims, fill_value=fill)
            if attrs_from is None:
                attrs = {}
            else:
                attrs = {a: attrs_from.getncattr(a)
                                for a in attrs_from.ncattrs()}
            if lod is not None:
                attrs['lod'] = lod
            if units is not None:
                attrs['units'] = units
            if attrs:
                v.setncatts(attrs)
            return v

        # Create dimensions
        fout.createDimension('time',  rt.nt     )
        fout.createDimension('z',     rt.nz     )
        fout.createDimension('zw',    rt.nz-1   )
        fout.createDimension('zsoil', rt.nz_soil)
        fout.createDimension('x',     rt.nx     )
        fout.createDimension('xu',    rt.nx-1   )
        fout.createDimension('y',     rt.ny     )
        fout.createDimension('yv',    rt.ny-1   )

        # Global attributes
        fout.setncatts({
            'origin_z': rt.origin_z,
            })

        # Create and write dimension variables
        mkvar('time',  ('time',), units=f'seconds since {rt.simulation.start_time}'
                )[:] = rt.times_sec
        mkvar('z',     ('z',)    )[:] = rt.z_levels[:]
        mkvar('zw',    ('zw',)   )[:] = rt.z_levels_stag[:]
        mkvar('zsoil', ('zsoil',))[:] = rt.z_soil_levels[:]
        mkvar('y',     ('y',)    )[:] = rt.dy/2 + rt.dy*np.arange(rt.ny,dtype=dtdefault)
        mkvar('x',     ('x',)    )[:] = rt.dx/2 + rt.dx*np.arange(rt.nx,dtype=dtdefault)
        mkvar('yv',    ('yv',)   )[:] = rt.dy*np.arange(1,rt.ny,dtype=dtdefault)
        mkvar('xu',    ('xu',)   )[:] = rt.dx*np.arange(1,rt.nx,dtype=dtdefault)

        # Create init variables
        mkvar('init_atmosphere_pt', ('z', 'y', 'x'),  cfg.output.lod.pt,  orig_lod=2)
        mkvar('init_atmosphere_qv', ('z', 'y', 'x'),  cfg.output.lod.qv,  orig_lod=2)
        mkvar('init_atmosphere_u',  ('z', 'y', 'xu'), cfg.output.lod.uvw, orig_lod=2)
        mkvar('init_atmosphere_v',  ('z', 'yv', 'x'), cfg.output.lod.uvw, orig_lod=2)
        mkvar('init_atmosphere_w',  ('zw', 'y', 'x'), cfg.output.lod.uvw, orig_lod=2)
        mkvar('init_soil_t',        ('zsoil', 'y', 'x'), 2)
        mkvar('init_soil_m',        ('zsoil', 'y', 'x'), 2)

        # Create forcing variables
        if not rt.nested_domain:
            # surface pressure (scalar)
            mkvar('surface_forcing_surface_pressure', ('time',))

            # boundary - vertical slices from left, right, south, north, top
            mkvar('ls_forcing_left_pt',  ('time', 'z',  'y' ), 2)
            mkvar('ls_forcing_right_pt', ('time', 'z',  'y' ), 2)
            mkvar('ls_forcing_south_pt', ('time', 'z',  'x' ), 2)
            mkvar('ls_forcing_north_pt', ('time', 'z',  'x' ), 2)
            mkvar('ls_forcing_top_pt',   ('time', 'y',  'x' ), 2)

            mkvar('ls_forcing_left_qv',  ('time', 'z',  'y' ), 2)
            mkvar('ls_forcing_right_qv', ('time', 'z',  'y' ), 2)
            mkvar('ls_forcing_south_qv', ('time', 'z',  'x' ), 2)
            mkvar('ls_forcing_north_qv', ('time', 'z',  'x' ), 2)
            mkvar('ls_forcing_top_qv',   ('time', 'y',  'x' ), 2)

            mkvar('ls_forcing_left_u',   ('time', 'z',  'y' ), 2)
            mkvar('ls_forcing_right_u',  ('time', 'z',  'y' ), 2)
            mkvar('ls_forcing_south_u',  ('time', 'z',  'xu'), 2)
            mkvar('ls_forcing_north_u',  ('time', 'z',  'xu'), 2)
            mkvar('ls_forcing_top_u',    ('time', 'y',  'xu'), 2)

            mkvar('ls_forcing_left_v',   ('time', 'z',  'yv'), 2)
            mkvar('ls_forcing_right_v',  ('time', 'z',  'yv'), 2)
            mkvar('ls_forcing_south_v',  ('time', 'z',  'x' ), 2)
            mkvar('ls_forcing_north_v',  ('time', 'z',  'x' ), 2)
            mkvar('ls_forcing_top_v',    ('time', 'yv', 'x' ), 2)

            mkvar('ls_forcing_left_w',   ('time', 'zw', 'y' ), 2)
            mkvar('ls_forcing_right_w',  ('time', 'zw', 'y' ), 2)
            mkvar('ls_forcing_south_w',  ('time', 'zw', 'x' ), 2)
            mkvar('ls_forcing_north_w',  ('time', 'zw', 'x' ), 2)
            mkvar('ls_forcing_top_w',    ('time', 'y',  'x' ), 2)

            # prepare influx/outflux area sizes
            zstag_all = np.r_[0., rt.z_levels_stag, rt.ztop]
            zwidths = zstag_all[1:] - zstag_all[:-1]
            verbose('Z widths: {}', zwidths)
            areas_xb = (zwidths * rt.dy)[:,ax_]
            areas_yb = (zwidths * rt.dx)[:,ax_]
            areas_zb = rt.dx * rt.dy

            tmask_left  = rt.obstacle_mask[:, :, 0]
            tmask_right = rt.obstacle_mask[:, :, rt.nx-1]
            tmask_south = rt.obstacle_mask[:, 0, :]
            tmask_north = rt.obstacle_mask[:, rt.ny-1, :]
            ntmask_left  = ~tmask_left
            ntmask_right = ~tmask_right
            ntmask_south = ~tmask_south
            ntmask_north = ~tmask_north

            area_boundaries = ((areas_xb * ntmask_left ).sum() +
                               (areas_xb * ntmask_right).sum() +
                               (areas_yb * ntmask_south).sum() +
                               (areas_yb * ntmask_north).sum() +
                                areas_zb*rt.nx*rt.ny)

        log('Writing values for initialization variables')
        with netCDF4.Dataset(rt.paths.intermediate.vinterp) as fin:
            fiv = fin.variables

            # geostrophic wind (1D)
            if cfg.output.geostrophic_wind:
                mkvar('ls_forcing_ug', ('time', 'z'))
                mkvar('ls_forcing_vg', ('time', 'z'))

            # write values for initialization variables
            fov['init_atmosphere_pt'][:] = lod2_to_lod(cfg.output.lod.pt,  fiv['init_atmosphere_pt'][0, :, :, :])
            fov['init_atmosphere_qv'][:] = lod2_to_lod(cfg.output.lod.qv,  fiv['init_atmosphere_qv'][0, :, :, :])
            fov['init_atmosphere_u'][:]  = lod2_to_lod(cfg.output.lod.uvw, fiv['init_atmosphere_u'][0, :, :, 1:]) #TODO fix stag
            fov['init_atmosphere_v'][:]  = lod2_to_lod(cfg.output.lod.uvw, fiv['init_atmosphere_v'][0, :, 1:, :]) #TODO fix stag
            fov['init_atmosphere_w'][:]  = lod2_to_lod(cfg.output.lod.uvw, fiv['init_atmosphere_w'][0, :, :, :] )
            fov['init_soil_t'][:,:,:] = fiv['init_soil_t'][0,:,:,:]
            fov['init_soil_m'][:,:,:] = (fiv['init_soil_m'][0,:,:,:]
                    * rt.soil_moisture_adjust[ax_,:,:])

            # Write values for time dependent values
            if not rt.nested_domain:
                for it in range(rt.nt):
                    verbose('Processing timestep {}', it)

                    # surface pressure: in PALM, surface pressure is actually level 0 (zb) pressure!!!
                    fov['surface_forcing_surface_pressure'][it] = fiv['palm_hydrostatic_pressure_stag'][it,0]

                    # boundary conditions
                    fov['ls_forcing_left_pt' ][it,:,:] = fiv['init_atmosphere_pt'][it, :, :, 0]
                    fov['ls_forcing_right_pt'][it,:,:] = fiv['init_atmosphere_pt'][it, :, :, rt.nx-1]
                    fov['ls_forcing_south_pt'][it,:,:] = fiv['init_atmosphere_pt'][it, :, 0, :]
                    fov['ls_forcing_north_pt'][it,:,:] = fiv['init_atmosphere_pt'][it, :, rt.ny-1, :]
                    fov['ls_forcing_top_pt'  ][it,:,:] = fiv['init_atmosphere_pt'][it, rt.nz-1, :, :]

                    fov['ls_forcing_left_qv' ][it,:,:] = fiv['init_atmosphere_qv'][it, :, :, 0]
                    fov['ls_forcing_right_qv'][it,:,:] = fiv['init_atmosphere_qv'][it, :, :, rt.nx-1]
                    fov['ls_forcing_south_qv'][it,:,:] = fiv['init_atmosphere_qv'][it, :, 0, :]
                    fov['ls_forcing_north_qv'][it,:,:] = fiv['init_atmosphere_qv'][it, :, rt.ny-1, :]
                    fov['ls_forcing_top_qv'  ][it,:,:] = fiv['init_atmosphere_qv'][it, rt.nz-1, :, :]

                    if self.is_anelastic:
                        verbose('Using anelastic approximation for mass balancing')

                        # Calculate air density profile
                        pres = fiv['palm_hydrostatic_pressure'][it,:]
                        pt_prof = fiv['init_atmosphere_pt'][it,:,:,:].mean(axis=(1,2))
                        rho_air = PalmPhysics.rho_air_ideal_gas(pres, pt_prof)
                        del pres, pt_prof

                        # Multiply boundary grid point areas with air density
                        bfact_x = areas_xb * rho_air[:,ax_]
                        bfact_y = areas_yb * rho_air[:,ax_]
                        bfact_z = areas_zb * rho_air[-1]
                        area_factors = ((bfact_x * ntmask_left ).sum() +
                                        (bfact_x * ntmask_right).sum() +
                                        (bfact_y * ntmask_south).sum() +
                                        (bfact_y * ntmask_north).sum() +
                                         bfact_z*rt.nx*rt.ny)

                        verbose('Air density profile (kg/m3): {}', rho_air)
                        bal_unit1 = 'kg/s'
                    else:
                        verbose('Using boussinesq approximation for mass balancing')

                        bfact_x = areas_xb
                        bfact_y = areas_yb
                        bfact_z = areas_zb
                        area_factors = area_boundaries
                        bal_unit1 = 'm3/s'

                    # Perform mass balancing for U, V, W
                    uxleft = fiv['init_atmosphere_u'][it, :, :, 0]
                    uxleft[tmask_left] = 0.
                    uxright = fiv['init_atmosphere_u'][it, :, :, rt.nx-1]
                    uxright[tmask_right] = 0.
                    vysouth = fiv['init_atmosphere_v'][it, :, 0, :]
                    vysouth[tmask_south] = 0.
                    vynorth = fiv['init_atmosphere_v'][it, :, rt.ny-1, :]
                    vynorth[tmask_north] = 0.
                    wztop = fiv['init_atmosphere_w'][it, rt.nz-2, :, :]#nzw=nz-1
                    mass_disbalance = (  (uxleft  * bfact_x).sum()
                                       - (uxright * bfact_x).sum()
                                       + (vysouth * bfact_y).sum()
                                       - (vynorth * bfact_y).sum()
                                       - (wztop   * bfact_z).sum() )
                    mass_corr_v = mass_disbalance / area_factors
                    log('Mass disbalance: {:8g} {} (correction = {:8g} m/s)',
                        mass_disbalance, bal_unit1, mass_corr_v)
                    uxleft[ntmask_left] -= mass_corr_v
                    uxright[ntmask_right] += mass_corr_v
                    vysouth[ntmask_south] -= mass_corr_v
                    vynorth[ntmask_north] += mass_corr_v
                    wztop += mass_corr_v

                    # Verify mass balance
                    if cfg.output.check_mass_balance and cfg.verbosity >= 1:
                        mass_disbalance = (  (uxleft  * bfact_x).sum()
                                           - (uxright * bfact_x).sum()
                                           + (vysouth * bfact_y).sum()
                                           - (vynorth * bfact_y).sum()
                                           - (wztop   * bfact_z).sum() )
                        mass_corr_v = mass_disbalance / area_factors
                        log('Mass balanced:   {:8g} {} (correction = {:8g} m/s)',
                            mass_disbalance, bal_unit1, mass_corr_v)

                    # Write U, V, W
                    fov['ls_forcing_left_u' ][it,:,:] = uxleft
                    fov['ls_forcing_right_u'][it,:,:] = uxright
                    fov['ls_forcing_south_u'][it,:,:] = fiv['init_atmosphere_u'][it, :, 0, 1:]
                    fov['ls_forcing_north_u'][it,:,:] = fiv['init_atmosphere_u'][it, :, rt.ny-1, 1:]
                    fov['ls_forcing_top_u'  ][it,:,:] = fiv['init_atmosphere_u'][it, rt.nz-1, :, 1:]

                    fov['ls_forcing_left_v' ][it,:,:] = fiv['init_atmosphere_v'][it, :, 1:, 0]
                    fov['ls_forcing_right_v'][it,:,:] = fiv['init_atmosphere_v'][it, :, 1:, rt.nx-1]
                    fov['ls_forcing_south_v'][it,:,:] = vysouth
                    fov['ls_forcing_north_v'][it,:,:] = vynorth
                    fov['ls_forcing_top_v'  ][it,:,:] = fiv['init_atmosphere_v'][it, rt.nz-1, 1:, :]

                    fov['ls_forcing_left_w' ][it,:,:] = fiv['init_atmosphere_w'][it, :, :, 0]
                    fov['ls_forcing_right_w'][it,:,:] = fiv['init_atmosphere_w'][it, :, :, rt.nx-1]
                    fov['ls_forcing_south_w'][it,:,:] = fiv['init_atmosphere_w'][it, :, 0, :]
                    fov['ls_forcing_north_w'][it,:,:] = fiv['init_atmosphere_w'][it, :, rt.ny-1, :]
                    fov['ls_forcing_top_w'  ][it,:,:] = wztop

                    # geostrophic wind (1D)
                    if cfg.output.geostrophic_wind:
                        fov['ls_forcing_ug'][it] = fiv['ls_forcing_ug'][it]
                        fov['ls_forcing_vg'][it] = fiv['ls_forcing_vg'][it]

            # Write chemical boundary conds
            if cfg.chem_species:
                log('Writing values for chemistry variables')

                convert_to_ppmv = set()
                for vn in cfg.chem_species:
                    vin = fiv[vn]

                    unit = vin.units
                    if (unit == cfg.chem_units.targets.kgm3
                            and not getattr(vin, 'non_gasphase', False)):
                        if not hasattr(vin, 'molar_mass'):
                            die('Variable {} needs to be converted to ppmv but it '
                                    'is missing molar mass!', vn)
                        convert_to_ppmv.add(vn)
                        unit = cfg.chem_units.targets.ppmv

                    mkvar('init_atmosphere_'+vn, ('z',), 1, unit, attrs_from=vin)
                    if not rt.nested_domain:
                        mkvar('ls_forcing_left_'+vn,  ('time','z','y'), 2, unit, attrs_from=vin)
                        mkvar('ls_forcing_right_'+vn, ('time','z','y'), 2, unit, attrs_from=vin)
                        mkvar('ls_forcing_south_'+vn, ('time','z','x'), 2, unit, attrs_from=vin)
                        mkvar('ls_forcing_north_'+vn, ('time','z','x'), 2, unit, attrs_from=vin)
                        mkvar('ls_forcing_top_'+vn,   ('time','y','x'), 2, unit, attrs_from=vin)

                # TODO move to separate plugin
                if cfg.postproc.nox_post_sum:
                    vin = fiv[cfg.postproc.nox_post_sum[0]]
                    unit = cfg.chem_units.targets.ppmv
                    mkvar('init_atmosphere_NOX', ('z',), 1, unit, attrs_from=vin)
                    if not rt.nested_domain:
                        mkvar('ls_forcing_left_NOX',  ('time','z','y'), 2, unit, attrs_from=vin)
                        mkvar('ls_forcing_right_NOX', ('time','z','y'), 2, unit, attrs_from=vin)
                        mkvar('ls_forcing_south_NOX', ('time','z','x'), 2, unit, attrs_from=vin)
                        mkvar('ls_forcing_north_NOX', ('time','z','x'), 2, unit, attrs_from=vin)
                        mkvar('ls_forcing_top_NOX',   ('time','y','x'), 2, unit, attrs_from=vin)

                if convert_to_ppmv:
                    log('Chemical quantities converted from kg/m3 to ppmv: {}', convert_to_ppmv)

                for it in range(1 if rt.nested_domain else rt.nt):
                    verbose('Processing timestep {}', it)

                    if convert_to_ppmv:
                        # calculate molar volume V/n = R*T/p
                        pres = fiv['palm_hydrostatic_pressure'][it,:][:,ax_,ax_]
                        t = fiv['init_atmosphere_pt'][it,:,:,:] * PalmPhysics.exner(pres)
                        mol_vol = t * PalmPhysics.R / pres # m3/mol

                    for vn in cfg.chem_species:
                        # Load timestep
                        v = fiv[vn][it,:,:,:]

                        if vn in convert_to_ppmv:
                            v *= mol_vol * (1e9 / fiv[vn].molar_mass) # kg/g*1e9 = 1e6

                        # PALM doesn't support 3D LOD=2 init for chem yet, we have
                        # to average the field
                        if it == 0:
                            fov['init_atmosphere_'+vn][:] = v.mean(axis=(1,2))

                        if not rt.nested_domain:
                            fov['ls_forcing_left_' +vn][it] = v[:,:,0]
                            fov['ls_forcing_right_'+vn][it] = v[:,:,-1]
                            fov['ls_forcing_south_'+vn][it] = v[:,0,:]
                            fov['ls_forcing_north_'+vn][it] = v[:,-1,:]
                            fov['ls_forcing_top_'  +vn][it] = v[-1,:,:]

                    if cfg.postproc.nox_post_sum:
                        if it == 0:
                            fov['init_atmosphere_NOX'][:] = sum(fov['init_atmosphere_'+vn][:]
                                                                for vn in cfg.postproc.nox_post_sum)
                        if not rt.nested_domain:
                            fov['ls_forcing_left_NOX'][it] = sum(fov['ls_forcing_left_'+vn][it]
                                                                for vn in cfg.postproc.nox_post_sum)
                            fov['ls_forcing_right_NOX'][it] = sum(fov['ls_forcing_right_'+vn][it]
                                                                for vn in cfg.postproc.nox_post_sum)
                            fov['ls_forcing_south_NOX'][it] = sum(fov['ls_forcing_south_'+vn][it]
                                                                for vn in cfg.postproc.nox_post_sum)
                            fov['ls_forcing_north_NOX'][it] = sum(fov['ls_forcing_north_'+vn][it]
                                                                for vn in cfg.postproc.nox_post_sum)
                            fov['ls_forcing_top_NOX'][it] = sum(fov['ls_forcing_top_'+vn][it]
                                                                for vn in cfg.postproc.nox_post_sum)

        if cfg.radiation.enabled:
            # Separate time dimension for radiation
            fout.createDimension('time_rad', rt.nt_rad)
            var = mkvar('time_rad', ('time_rad',))
            var[:] = rt.times_rad_sec

            # radiation variables
            var = mkvar('rad_sw_in', ('time_rad',), 1, 'W/m2')
            var[:] = rt.rad_swdown

            var = mkvar('rad_lw_in', ('time_rad',), 1, 'W/m2')
            var[:] = rt.rad_lwdown

            if rt.has_rad_diffuse:
                var = mkvar('rad_sw_in_dif', ('time_rad',), 1, 'W/m2')
                var[:] = rt.rad_swdiff

        log('Dynamic driver written successfully.')
