# -*- coding: utf-8 -*-
"""
Created on Tue Oct 31 16:12:55 2023

@author: ZMK
"""

import flopy
import flopy.utils.binaryfile as bf
import csv
import Base as base
import os
import json
import ModelPeriod
import numpy as np



def get_model_json(model_name):
      period_json=""
      prediction_path = base.model_dir + model_name +"\\prediction.json"
      with open(prediction_path,encoding='utf-8') as f:
             period_json = json.load(f)  
     
      return period_json; 
  
def get_model_period(model_name):
     period_json=""
     prediction_path = base.model_dir + model_name +"\\prediction.json"
     with open(prediction_path,encoding='utf-8') as f:
             period_json = json.load(f)  
     
     start_time = period_json["start_time"]
     end_time = period_json["end_time"]
      
     months = ModelPeriod.get_months_in_range_ym(start_time, end_time)
     return months; 
    

#观测井chart
def obsChartdata(model_name, row, column):
    
      row = int(row)-1
      column = int(column)-1
      dir = base.model_dir + model_name + "\\modflow.head"

      head = bf.HeadFile(dir)
      alldata = head.get_alldata()
      period = len(alldata)

      layer = 3

      xdata = []
      ydata = []
      result = {}
      for per in range(period):
           for lay in range(layer):
               if per % 3 == 0 and lay == 0:
                   md = (int)(lay / 3 + 1)
                   per_array = alldata[per][lay]

                   cell_data = (float)(per_array[row][column])
                   ydata.append(cell_data)
    
      period_json= get_model_json(model_name)
      
      start_time = period_json["start_time"]
      end_time = period_json["end_time"]
      
      months = ModelPeriod.get_months_in_range_ym(start_time, end_time)

      result = {"y_data": ydata, "x_data": months}
      return result

def getRowCloumnById(index_id):
    row = 104
    column =114
    count=0
    
    for  i in range(row):
        for j in range(column):
            if index_id == count:
                return (i,j)
            count = count +1
    return ""        
            
    
    
#地下水信息
def earthWaterChart(model_name, index_id):
    
      row_column =  getRowCloumnById(index_id)
      
      row = row_column[0]
      column = row_column[1]
      dir = base.model_dir + model_name + "\\modflow.head"

      head = bf.HeadFile(dir)
      alldata = head.get_alldata()
      period = len(alldata)

      layer = 3

      ydata = []
      result = {}
      for per in range(period):
           for lay in range(layer):
               if per % 3 == 0 and lay == 0:
                  
                   per_array = alldata[per][lay]

                   cell_data = (float)(per_array[row][column])
                   ydata.append(cell_data)
    
      period_json= get_model_json(model_name)
      
      start_time = period_json["start_time"]
      end_time = period_json["end_time"]
      
      months = ModelPeriod.get_months_in_range_ym(start_time, end_time)

      result = {"y_data": ydata, "x_data": months}
      return result

def heatmapdata(model_name,period):
    dir = base.model_dir + model_name + "\\modflow.head"
      
    head = bf.HeadFile(dir)

    alldata = head.get_alldata()
    
    index = int(period)*3
    return alldata[index][0]


#水均衡计算
def waterEqu(model_name):
    if model_name == '202001_202212':  
        water_equ_path = base.prefix + "\\water_equ.json"
        with open(water_equ_path,encoding='utf-8') as f:
             data = json.load(f)
             return data       
    else:
        year = model_name[0:4]
        title =[year]
        dict ={"title":title}
        
        celldata = np.array(base.water_equ2022).tolist()
      
        predict_json= get_model_json(model_name)
        
        a1=float(celldata[0])
        a2=float(celldata[1])
        a3=float(celldata[2])
        a4=float(celldata[3])
        
        b1=float(celldata[4])
        b2=float(celldata[5])
        b3=float(celldata[6])
        
        if predict_json["rain"]:
            a1= float(predict_json["rain"]["ratio"]) * float(celldata[0])  
            a3= float(predict_json["rain"]["ratio"]) * float(celldata[2]) 
            a4= float(predict_json["rain"]["ratio"]) * float(celldata[3])        
            b2= float(predict_json["rain"]["ratio"]) * float(celldata[5]) 
            b3= float(predict_json["rain"]["ratio"]) * float(celldata[6])     
        if predict_json["river"]:
            a2=  float(predict_json["river"]["ratio"]) * float(celldata[1])
            
        if predict_json["mine"]:
            b1=  b1    
        
        in_data= a1+a2+a3+a4
        out_data= b1 +b2 + b3
        float_data=[a1,a2,a3,a4,in_data,b1,b2,b3,out_data,in_data-out_data]
        
        inarray=[]
        inarray.append({"name":"降水入渗量","value":a1})
        inarray.append({"name":"河流入渗量","value":a2})
        inarray.append({"name":"L1侧向补给量","value":a3})
        inarray.append({"name":"L3侧向补给量","value":a4})
        outarray=[]
        outarray.append({"name":"人工开采量","value":b1})
        outarray.append({"name":"L1侧向流出量","value":b2})
        outarray.append({"name":"L3侧向流出量","value":b3})
        pie1={str(year):inarray}
        pie2={str(year):outarray}
        
        dict["pie1"]=pie1
        dict["pie2"]=pie2
        
        array2d=[]
        array2d.append([str(year)])
        for i in range(len(float_data)):
            tmp=[]
            tmp.append(str(float_data[i]))
            array2d.append(tmp)
        dict["data"]=array2d
        return dict
        

#导出csv文件
def exportCsV(model_name):
    
    dir = base.model_dir + model_name + "\\modflow.head"
    out_path = base.model_dir + model_name + "\\output\\"
    if not os.path.exists(out_path):
         os.mkdir(out_path)  
    
    head = bf.HeadFile(dir)

    alldata = head.get_alldata()
    month = len(alldata)
    layer = 3

    for i in range(month):
       for j in range(layer):
          if i % 3 == 0:
            md = (int)(i / 3 + 1)
            filename = out_path + str(md) + '-' + str(j+1) + '.csv'
            f = open(filename, 'w', newline='')
            writer = csv.writer(f)
            for p in alldata[i][j]:
               writer.writerow(p)
            f.close()

    return out_path