#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Oct 30 14:46:00 2023

@author: tiagomandu
"""

import numpy as np
import math
import glob
import os

"""
 AUTHOR @ JEAN-FRANCOIS RIBAUD

 DATE: OCTOBER 2018

 PURPOSE: PERFORM THE STEINER ET AL 1995 CONVECTIVE-STRATIFORM CLASSIFICATION
          FROM CARTESIAN GRID AT 2 OR 3 KM HEIGHT AND BASED ON CORRECTED HORIZONTAL
          REFLECTIVITY.

          MORE INFORMATION ABOUT THE METHODOLOGY IN STEINER ET AL 1995

"""

#*****************************
PathName = '/oper/share/radar/sdcsc/chapeco/cappi/cappi3km_bin/2019/05/'
path_output = '/share/apgmet/dist/Tiago/cs_chapeco/2019/05/'

#*************************************************************************


#########################################
def calc_critical_delta_dbz(dbz):
    """
    """
    if dbz < 0:
        return 10
    elif dbz >= 42.43:
        return 0
    else:
        return 10 - (dbz**2) / 180


#########################################
def calc_convective_radius(dbz):
    """
    return convective radius in [km]
    """
    if dbz < 25:
        return 1
    elif dbz < 30:
        return 2
    elif dbz < 35:
        return 3
    elif dbz < 40:
        return 4
    else:
        return 5


#############################
def calc_background_dbz_avg(Z,x,y,r,res):
    """
    Z = reflectivity [mm6.m-3]
    return ZH[dBZ]
    """

    # search square
    x_inf = 0 if x-(r/res) < 0 else x-(r/res)
    x_sup = (len(Z)-1) if x+(r/res) > (len(Z)-1) else x+(r/res)
    y_inf = 0 if y-(r/res) < 0 else y-(r/res)
    y_sup = (len(Z)-1) if y+(r/res) > (len(Z)-1) else y+(r/res)

    x_inf = int(x_inf)
    y_inf = int(y_inf)
    x_sup = int(x_sup)
    y_sup = int(y_sup)
    
    # slice we care about
    #print(x_inf, x_sup, y_inf, y_sup)
    sub_Z = Z[x_inf:x_sup, y_inf:y_sup]
    ysub_grid, xsub_grid = np.meshgrid(np.arange(y_sup-y_inf), np.arange(x_sup-x_inf))

    # Make a circular mask
    distance = ((xsub_grid - x)/res)**2 + ((ysub_grid - y)/res)**2
    mask = distance > r**2
    area = np.ma.masked_where(mask, sub_Z)
    #if all values were masked, do nothing
    if (np.ma.count(area) > 0):
        average = np.nanmean(area) #[mm6.m-3]
        # Return average dBZ
        if average > 0:
            return 10 * math.log(average, 10)
        else: # here we considerer initial negative ZH - be careful
            return 0 
    else:
        return 0


###############################
def mark_convective_radius(ZH,CS,x,y,r,res):
    """
    """
    
    # search square
    x_inf = 0 if x-(r/res) < 0 else x-(r/res)-1
    x_sup = (len(ZH)-1) if x+(r/res) > (len(ZH)-1) else x+(r/res)+1
    y_inf = 0 if y-(r/res) < 0 else y-(r/res)-1
    y_sup = (len(ZH)-1) if y+(r/res) > (len(ZH)-1) else y+(r/res)+1

    for n in range(int(x_inf),int(x_sup)):
        for m in range(int(y_inf),int(y_sup)):
            # See if distance is <= convective radius
            dist_2d = (((n-x)/res)**2 + ((m-y)/res)**2)**0.5
            if dist_2d <= r:
                # Mark convective
                if CS[n,m] != 2:
                    CS[n,m] = 1
                            
    return CS




###########################################################################################



    ################
    # STEINER ALGO #
    ################

#   0 = stratiform
#   1 = convective
#   2 = convective center
#   3 = invalid values/non-calculated
#   99 = calculating


#*****************************
# INPUTS

grid_dim = 500 # [km] # diametro (total)
res = 1. # [km]

# background radius from Steiner 1995
radius = 11 # [km] 

files = sorted(glob.glob(PathName+'R12132246_*raw'))
total = len(files)

for p in range(len(files)):

    file = files[p]
    print(file)

    #************
    # LOAD DATA
    with open(file,'rb') as f:
        try:
            ZH = np.fromfile(f, dtype=np.float32)
            ZH = np.reshape(ZH, [grid_dim,grid_dim]) # [dBZ]
        except:
            print(file,' DO NOT EXIST - PROBLEM !!!')
    ZH[ZH==-99] = np.nan 

    #****************************************************
    # TO CALCULATE AVERAGE OF Z (INSTEAD OF ZH[dBZ] !!!)
    Z = 10**(np.copy(ZH)/10)  # [mm6.m-3] 

    #*********************
    # DEFINE OUTPUT GRID 
    CS = np.zeros([grid_dim,grid_dim], dtype=np.int8) + 3 # Convective-Stratiform matrix


    #*************
    for i in range(grid_dim):
        for j in range(grid_dim):
            if ZH[i,j] >= 40 : # STEP 1: Mark any grid >= 40 dBZ as a convective center
                CS[i,j] = 2
            else: # STEP 2: calculate background average for point
                background_dbz_avg = calc_background_dbz_avg(Z,i,j,radius,res)
                critical_delta = calc_critical_delta_dbz(background_dbz_avg) 

                # If echo-background delta is greater than critical value, mark convective 
                if (ZH[i,j] - background_dbz_avg) > critical_delta:
                    CS[i,j] = 2
                else:
                    CS[i,j] = 99 # Mark for later calculation


    #*************
    for i in range(grid_dim):
        for j in range(grid_dim):
            if CS[i,j] == 2 :
                background_dbz_avg = calc_background_dbz_avg(Z,i, j, radius,res)
                convective_radius = calc_convective_radius(background_dbz_avg)
                mark_convective_radius(ZH, CS, i, j, convective_radius,res)


    #*************
    for i in range(grid_dim):
        for j in range(grid_dim):
            if CS[i,j] == 99: # Set any uncalced values to 3 - invalid
                CS[i,j] = 3
            if CS[i,j] == 2: #Set Convective Cores to be convection 
                CS[i,j] = 1


    #*************
    # SAVE
    
    np.savez_compressed(path_output + os.path.basename(file) + '_chapeco_CS.npz', CS)
