# -*- coding: utf-8 -*-

# @File    : new_cangnan.py
# @Date    : 2023-07-25
# @Author  : Dovelet
# -*- coding: utf-8 -*-

import netCDF4
import numpy as np
import os
import datetime
import baysalt.strDateTime as strDateTime
from scipy.interpolate import interp2d
from scipy.interpolate import interpn
import scipy.integrate
import json
import math

m_hours = list(range(216))
fvcom_dir='/data/Output_147/FVCOM_ECS_storm/Run/'

class Cangnan_Single:
    def __init__(self, _lon, _lat, _input, _output_txt, _output_json, _name, **_new_dict):
        self._lon = _lon
        self._lat = _lat
        self._input = _input
        self._output_txt = _output_txt
        self._output_json = _output_json
        self._name = _name
        self._new_dict = _new_dict
        self.p=None

    def getnum(self, _date, _idx):
        """
        获取ww3_djk数据
        :param _date: 任务日期
        :param _idx: 小时
        :return: 有效波高，平均周期，谱峰周期，平均波向
        """
        # 时间处理，将UTC转为CST
        date1 = datetime.datetime.strptime(f'{_date}00', '%Y%m%d%H')
        _idx = _idx - 8  # 向前8个钟头，找到CST当天的0点
        date2 = date1 + datetime.timedelta(hours=_idx)
        idx = _idx % 24
        filename = self._input + date2.strftime('%Y%m%d') + '/ww3.' + date2.strftime(
            '%Y%m%d') + '.nc'  # ww3 nc文件名

        w = netCDF4.Dataset(filename)  # 读取ww3_Wp的nc文件，传入w中
        lons = w.variables['longitude'][:]
        lats = w.variables['latitude'][:]
        ww3hs = w.variables['hs'][idx, :, :]  # 有效波高
        ww3t02 = w.variables['t02'][idx, :, :]  # 平均周期
        ww3fp = w.variables['fp'][idx, :, :]  # 谱峰频率
        ww3dir = w.variables['dir'][idx, :, :]  # 波向
        ww3swellhs = w.variables['phs1'][idx, :, :]  # 涌浪有效波高
        ww3swelltm = w.variables['ptp1'][idx, :, :]  # 涌浪谱峰周期
        ww3lm = w.variables['lm'][idx, :, :]  # 波长

        fhs = interp2d(lons, lats, ww3hs, kind='linear')
        ft02 = interp2d(lons, lats, ww3t02, kind='linear')
        ffp = interp2d(lons, lats, ww3fp, kind='linear')
        fdir = interp2d(lons, lats, ww3dir, kind='linear')
        fswellhs = interp2d(lons, lats, ww3swellhs, kind='linear')
        fswelltm = interp2d(lons, lats, ww3swelltm, kind='linear')
        flm = interp2d(lons, lats, ww3lm, kind='linear')

        hs = fhs(self._lon, self._lat)
        # H4
        t02 = ft02(self._lon, self._lat)
        # H10
        mwd = fdir(self._lon, self._lat)
        # pp
        PP = 1 / ffp(self._lon, self._lat)
        # phs1
        phs1 = fswellhs(self._lon, self._lat)
        # ptp1
        ptp1 = fswelltm(self._lon, self._lat)
        # LM
        lm = flm(self._lon, self._lat)

        if phs1[0] >= 1e+10 or ptp1[0] >= 1e+10:  # 涌浪波高或谱峰周期为空时用一维谱计算
            ef_filename = self._input + date2.strftime('%Y%m%d') + '/ww3.' + date2.strftime(
                '%Y%m%d') + '_ef.nc'  # ww3_ef nc文件名
            ww3_ef = netCDF4.Dataset(ef_filename)
            ww3ef = ww3_ef.variables['ef'][idx, :, :, :]
            f = ww3_ef.variables['f'][:]

            Z, Y, X = np.meshgrid(f, self._lat, self._lon, indexing='ij')
            ef = interpn((f, lats, lons), ww3ef, (Z, Y, X), method='linear')  # 维度和值的维度是正的

            # 涌浪要素
            # 新的涌浪计算公式
            sf1 = ef[:18, 0, 0]
            fs1 = f[:18]
            phs1[0] = 4 * math.sqrt(scipy.integrate.trapz(sf1, fs1))  # 计算涌浪有效波高
            ptp1[0] = math.sqrt(
                scipy.integrate.trapz(sf1, fs1) / scipy.integrate.trapz(np.multiply(fs1 ** 2, sf1), fs1))  # 计算涌浪有效周期

        hs_rtn = np.round(hs[0], 2)
        t02_rtn = np.round(t02[0], 2)
        dir_rtn = np.round(mwd[0], 2)
        PP_rtn = np.round(PP[0], 2)
        phs1_rtn = np.round(phs1[0], 2)
        ptp1_rtn = np.round(ptp1[0], 2)
        lm_rtn = np.round(lm[0], 2)

        return hs_rtn, t02_rtn, dir_rtn, PP_rtn, phs1_rtn, ptp1_rtn, lm_rtn

    def getfore(self, _date, _idx):
        """
        获取ww3_djk数据,流场数据
        :param _date: 任务日期
        :param _idx: 小时
        :return: 有效波高，平均周期，谱峰周期，平均波向
        """
        hs_rtn, t02_rtn, dir_rtn, PP_rtn, phs1_rtn, ptp1_rtn, lm_rtn = self.getnum(_date, _idx)

        fh_css,fh_csm,fh_csb,fh_cds,fh_cdm,fh_cdb=self.getcurrent(_date,_idx,self.p)
        return str(format(hs_rtn, '.2f')).rjust(12, " "), \
               str(format(t02_rtn, '.2f')).rjust(12, " "), \
               str(format(dir_rtn, '.2f')).rjust(12, " "), \
               str(format(PP_rtn, '.2f')).rjust(12, " "), \
               str(format(phs1_rtn, '.2f')).rjust(12, " "), \
               str(format(ptp1_rtn, '.2f')).rjust(12, " "), \
               str(format(lm_rtn, '.2f')).rjust(12, " "), \
                  str(format(fh_css, '.2f')).rjust(12, " "), \
                  str(format(fh_csm, '.2f')).rjust(12, " "), \
                    str(format(fh_csb, '.2f')).rjust(12, " "), \
                    str(format(fh_cds, '.2f')).rjust(12, " "), \
                    str(format(fh_cdm, '.2f')).rjust(12, " "), \
                    str(format(fh_cdb, '.2f')).rjust(12, " ")



    def getcurrent(self,_date,_idx,p):
        """
        获取fvcom数据
        :param _date: 日期
        :param _idx: 预报小时（需要做时间处理）
        :param p: fvcom网格点的索引
        :return: "currentspeeds","currentspeedm","currentspeedb","currentdirs","currentdirm","currentdirb"
        """
        # 时间处理，将UTC转为CST
        if _idx>=120:
            currentspeeds, currentspeedm, currentspeedb, currentdirs, currentdirm, currentdirb=np.nan,np.nan,np.nan,np.nan,np.nan,np.nan
            return currentspeeds, currentspeedm, currentspeedb, currentdirs, currentdirm, currentdirb
        date1 = datetime.datetime.strptime(f'{_date}00', '%Y%m%d%H')
        _idx = _idx - 8  # 向前8个钟头，找到CST当天的0点
        date2 = date1 + datetime.timedelta(hours=_idx)
        idx = _idx % 24
        if _idx<0:
            filename=os.path.join(fvcom_dir,date2.strftime('%Y%m%d')+'/output'+'/forecast_0001.nc')
        else:

            # 一个小天才的写法，哈哈哈哈
            filename=os.path.join(fvcom_dir,date2.strftime('%Y%m%d')+'/output'+'/forecast_0001.nc')
            if os.path.exists(filename)==False:
                filename=os.path.join(fvcom_dir,date1.strftime('%Y%m%d')+'/output'+'/forecast_'+str(int(_idx/24)+1).zfill(4)+'.nc')
        f_c = netCDF4.Dataset(filename) #fvcom nc
        u = f_c.variables['u'][:]
        v = f_c.variables['v'][:]
        currentspeeds = np.sqrt(u[idx,0, p] ** 2 + v[idx,0, p] ** 2)
        currentdirs= np.arctan2(v[idx,0, p], u[idx,0, p]) * 180 / np.pi
        currentspeedm = np.sqrt(u[idx,7, p] ** 2 + v[idx,7, p] ** 2)
        currentdirm = np.arctan2(v[idx,1, p], u[idx,1, p]) * 180 / np.pi
        currentspeedb = np.sqrt(u[idx,14, p] ** 2 + v[idx,14, p] ** 2)
        currentdirb = np.arctan2(v[idx,14, p], u[idx,14, p]) * 180 / np.pi

        currentspeeds=np.round(currentspeeds,2)
        currentspeedm=np.round(currentspeedm,2)
        currentspeedb=np.round(currentspeedb,2)
        currentdirs=np.round(currentdirs,2)
        currentdirm=np.round(currentdirm,2)
        currentdirb=np.round(currentdirb,2)

        return currentspeeds,currentspeedm,currentspeedb,currentdirs,currentdirm,currentdirb


    def get_current_sp(self,_date):
        """
        找到最近的fvcom数据点的索引
        :param _date: 日期。
        :return: 点位索引
        """
        _date = str(_date)
        filename = os.path.join(fvcom_dir, _date + '/output' + '/forecast_0001.nc')
        f_c = netCDF4.Dataset(filename)  # fvcom nc
        lonc= f_c.variables['lonc'][:]
        latc= f_c.variables['latc'][:]
        D=np.sqrt((lonc-self._lon)**2+(latc-self._lat)**2)
        idx=np.argmin(D)

        return idx

    @staticmethod
    def getTime(_date, _hour):
        """
        获取预报的日期/小时
        :param _date: 任务日期
        :param _hour: 小时
        :return: 日期/小时(10位长度的字符串)
        """
        if _hour == 215:
            _hour = 215

        date = strDateTime.convertToTime(_date)
        date_t = date + datetime.timedelta(hours=int(_hour))

        return date_t.strftime('%Y%m%d%H')

    @staticmethod
    def nan2str(_data):
        # 将nan转换为字符串
        if np.isnan(_data):
            return 'NaN'
        else:
            return _data

    @staticmethod
    def getJsonData(date, fc_hs, fc_dir, fc_t01, fc_PP, fc_phs1, fc_ptp1, fc_lm, fc_css, fc_csm, fc_csb, fc_cds, fc_cdm,fc_cdb):
        # sourcery skip: dict-literal, inline-immediately-returned-variable, merge-dict-assign
        postData = dict()
        postData['date'] = date
        postData['swh'] = fc_hs
        postData['mwd'] = fc_dir
        postData['mwp'] = fc_t01
        postData['pp1d'] = fc_PP
        postData['shts'] = fc_phs1
        postData['mpts'] = fc_ptp1
        postData['mwl'] = fc_lm
        postData['currentspeeds'] = Cangnan_Single.nan2str(fc_css)
        postData['currentspeedm'] = Cangnan_Single.nan2str(fc_csm)
        postData['currentspeedb'] = Cangnan_Single.nan2str(fc_csb)
        postData['currentdirs'] = Cangnan_Single.nan2str(fc_cds)
        postData['currentdirm'] = Cangnan_Single.nan2str(fc_cdm)
        postData['currentdirb'] = Cangnan_Single.nan2str(fc_cdb)
        return postData



    def grttxt(self, _date):
        """
        构建预报单信息
        :param _date: 预报单日期
        :return: 0-生成预报数据成功  大于等于1-生成预报数据失败
        """

        today = strDateTime.getToday()
        file_path = self._output_txt
        if not os.path.exists(file_path):
            os.makedirs(file_path)

        # 写入文件内容----------------------------------begin
        _date0 = _date
        list_hours = [m_hours]
        # 写入文件头----------------------------------begin
        fh_time = "time".rjust(12, " ")  # 第一列日期留空白
        fh_hs = "swh".rjust(12, " ")  # 有效波高
        fh_dir = "mwd".rjust(12, " ")  # 波向
        fh_t01 = "mwp".rjust(12, " ")  # 平均周期
        fh_lm = 'mwl'.rjust(12, " ")  # 波长
        fh_PP = "pp1d".rjust(12, " ")  # 谱峰周期
        fh_phs1 = 'shts'.rjust(12, " ")  # 涌浪有效波高
        fh_ptp1 = 'mpts'.rjust(12, " ")  # 涌浪谱峰周期
        # 流场数据开头
        fh_css='currentspeeds'.rjust(12, " ")  # 表面流速
        fh_csm='currentspeedm'.rjust(12, " ")  # 中层流速
        fh_csb='currentspeedb'.rjust(12, " ")  # 底层流速
        fh_cds='currentdirs'.rjust(12, " ")  # 表面流向
        fh_cdm='currentdirm'.rjust(12, " ")  # 中层流向
        fh_cdb='currentdirb'.rjust(12, " ")  # 底层流向

        self.p = self.get_current_sp(_date)

        for k in range(1):
            tmp = list_hours[k]
            file_name = self._name + strDateTime.addDays(_date, k) + ".txt"
            file = os.path.join(file_path, file_name)

            with open(file, "w") as fl:
                fl.writelines([fh_time, fh_hs, fh_dir, fh_t01, fh_lm, fh_PP, fh_phs1, fh_ptp1,fh_css,fh_csm,fh_csb,fh_cds,fh_cdm,fh_cdb,"\n"])
                # 写入文件头----------------------------------end

                for i in range(max(tmp) + 1):  # +72

                    fc_time = self.getTime(_date0, i)  # 第一列日期留空白

                    fc_hs, fc_t01, fc_dir, fc_PP, fc_phs1, fc_ptp1, fc_lm,fc_css,fc_csm,fc_csb,fc_cds,fc_cdm,fc_cdb = self.getfore(_date0, i)  # 获得数据

                    if fc_PP.strip() == "--":  # 谱峰周期为空时，采用2.99代替
                        fc_PP = str(2.99).rjust(12, " ")
                    if i in m_hours:
                        fl.writelines([fc_time, fc_hs, fc_dir, fc_t01, fc_lm, fc_PP, fc_phs1, fc_ptp1, fc_css,fc_csm,fc_csb,fc_cds,fc_cdm,fc_cdb,"\n"])
                        print(f"hour in={str(i)}")
                    else:
                        print(f"hour={str(i)}")

    def grtjson(self, _date):

        file_path = self._output_json
        if not os.path.exists(file_path):
            os.makedirs(file_path)

        _date0 = _date
        list_hours = [m_hours]
        self.p = self.get_current_sp(_date)
        # 创建json数据----------------------------------begin
        for k in range(1):
            tmp = list_hours[k]
            list = []
            for i in range(max(tmp) + 1):

                fc_time = self.getTime(_date0, i)  # 第一列日期留空白

                fc_hs, fc_t01, fc_dir, fc_PP, fc_phs1, fc_ptp1, fc_lm = self.getnum(_date0, i)  # 获取数据
                fc_css, fc_csm, fc_csb, fc_cds, fc_cdm, fc_cdb=self.getcurrent(_date0,i,self.p)

                if np.isnan(fc_PP):  # 谱峰周期为空时，采用2.99代替
                    fc_PP = 2.99

                JsonData = self.getJsonData(fc_time, fc_hs, fc_dir, fc_t01, fc_lm, fc_PP, fc_phs1, fc_ptp1,fc_css,fc_csm,fc_csb,fc_cds,fc_cdm,fc_cdb)  # dict()
                list.append(JsonData)
                # 创建json数据----------------------------------end
                # 创建json文件----------------------------------begin
                file_name = self._name + strDateTime.addDays(_date, k) + ".json"
                file = os.path.join(file_path, file_name)
                # new_dict = {"location": {"No": "rzlssh01", "name": "日照岚山港区原油码头", "lon": lon_g, "lat": lat_g}, "data": list}
                new_dict = {"location": self._new_dict, "data": list}
                with open(file, "w") as f:
                    json.dump(new_dict, f, ensure_ascii=False, indent=10)
                # 写入文件内容----------------------------------end


if __name__ == '__main__':
    input_path = '/data/ForecastSystem/WW3_6.07/Output/ww3_WPacific/'
    py_file_path = os.path.dirname(os.path.abspath(__file__))
    output_path_txt = os.path.join(py_file_path, './Data/single_txt/')
    output_path_json = os.path.join(py_file_path, './Data/single_json/')
    if not os.path.exists(output_path_txt):
        os.makedirs(output_path_txt)
    if not os.path.exists(output_path_json):
        os.makedirs(output_path_json)

    location_cnhd1 = {"No": "cnhd1", "name": "苍南核电1", "lon": 120.600, "lat": 27.260}
    name_cnhd1 = "cnhd1_120.521E_27.201N_"
    c=Cangnan_Single(120.600, 27.260, input_path, output_path_txt, output_path_json, name_cnhd1, **location_cnhd1)
    #c.grttxt('20230721')
    print('json')
    c.grtjson('20230721')
    print('end')
