#!/usr/bin/env python3
# Copyright 2023 - Barcelona Supercomputing Center
# Authors: Amirpasha Mozaffari
# MIT License

import yaml
from yaml.loader import SafeLoader
import os 
import argparse
import sys

from DataCouplerUtils import read_DataCoupler_yaml_input
 
# It will readin the YAML file that is generated by the RunExample and generate a namcouple list before 


def append_multiple_lines_to_file(file_path, lines_to_append):
    with open(file_path, 'a') as file:
       # for line in lines_to_append:
        file.write(lines_to_append + '\n')


def add_field_to_namcouple(send_id, receive_id, timestepsec, send_grid, receive_grid):
    # TODO add option for transformation and restart_file
    if send_grid == receive_grid:
        transfo = "LOCTRANS\n   AVERAGE"
    else:
        transfo = "SCRIPR\n   GAUSWGT LR SCALAR LATITUDE 1 9 2.0"

    lines_to_append =   """
  {send_id} {receive_id} 1 {timestepsec} 1 {restart_file} EXPOUT
  {send_grid} {receive_grid} LAG=0
  P  0  P  0
  {transfo}"""

    return lines_to_append.format(send_id=send_id, receive_id=receive_id, timestepsec=timestepsec, restart_file='rstas.nc', send_grid=send_grid, receive_grid=receive_grid, transfo=transfo)


def create_empty_file(filename):
    try:
        with open(filename, 'w') as file:
            pass  # The file is created but remains empty
        print(f"File '{filename}' created successfully.")
    except Exception as e:
        print(f"An error occurred: {e}")


def add_begin_to_namcouple(nfields, runtime, ModelNameSend, ModelNameReceive):
    lines_to_append = """# =================================================================================================
# General OASIS configuration
# =================================================================================================
 $NFIELDS
    {nfields}
 $END
# -------------------------------------------------------------------------------------------------
 $RUNTIME
    {runtime}
 $END
# -------------------------------------------------------------------------------------------------
 $NLOGPRT
 0 0
 $END
# -------------------------------------------------------------------------------------------------
 $STRINGS 

# =================================================================================================
# Fields send from {ModelNameSend} to {ModelNameReceive}
# ================================================================================================="""
    
    return(lines_to_append.format(nfields=nfields, runtime=runtime, ModelNameSend=ModelNameSend, ModelNameReceive=ModelNameReceive))


def add_end_to_namecouple():
    lines_to_append = """# -------------------------------------------------------------------------------------------------
$END
# ================================================================================================="""

    return lines_to_append


def create_yaml_conf_name_ifs_toy_model(filename,ModelName,run_length_sec,coupling_interval,nx,OasisOutputVars_readin):

    coupling_out_vars=[]
    
    coupling_in_vars = [] 
    for v in OasisOutputVars_readin.keys() :
        coupling_in_vars.append(OasisOutputVars_readin[v]['receive_id'])
    
    create_empty_file(filename)
    lines_to_append =""" # Copyright 2023 - Barcelona Supercomputing Center
# Author: Amirpasha Mozaffari
# TBD
---

# yaml metadata
model:
  ModelName     : {ModelName}

simulation:
  coupling_interval : {coupling_interval}
  run_length_sec : {run_length_sec}
  # Definition of the local partition grid cell 
  # L080_NX = 35718  # number of LAND grid cells at T159 resolution
  # L128_NX = 88838  # number of LAND grid cells at T255 resolution
  # L256_NX = 348528  # number of LAND grid cells at T511 resolution
  nx : {nx}

coupling: 
  out_vars : {out_vars}
  in_vars : {in_vars}
  restart_file : /path/ 
    """
    append_multiple_lines_to_file(filename, lines_to_append.format(ModelName=ModelName,coupling_interval=coupling_interval,run_length_sec=run_length_sec,nx=nx,out_vars=coupling_out_vars,in_vars=coupling_in_vars ))


if __name__ == '__main__':
    
    parser=argparse.ArgumentParser()
    parser.add_argument("--conf",type=str)
    parser.add_argument("--namcouple_output",type=str)
    parser.add_argument("--toy_model_yaml_conf_output",type=str)
    args = parser.parse_args()
    yaml_conf_file = args.conf
    namcouple_output =  args.namcouple_output
    toy_model_yaml_conf_output =  args.toy_model_yaml_conf_output 
 

    # Read the yaml file 

    ModelNameSend, ModelNameReceive, LogFileName, RunLengthSec, TimeStepSec, StartYear, StartMonth,   \
    StartDay, FixYear, GridInfo, LDebug, fileIn_vars_config, oasisOut_vars_config, OasisOutputVars_readin = read_DataCoupler_yaml_input(yaml_conf_file)

    number_of_fields = len(OasisOutputVars_readin)

    # Creating the namcouple file amd insert the read-in data from YAML file 
    create_empty_file(namcouple_output)
    lines=add_begin_to_namcouple(number_of_fields, RunLengthSec, ModelNameSend, ModelNameReceive)
    append_multiple_lines_to_file(namcouple_output, lines)
    
    for v in OasisOutputVars_readin.keys() :
        lines = add_field_to_namcouple(OasisOutputVars_readin[v]['send_id'], OasisOutputVars_readin[v]['receive_id'], \
                                    TimeStepSec, OasisOutputVars_readin[v]['send_grid_name'], OasisOutputVars_readin[v]['receive_grid_name'])
        append_multiple_lines_to_file(namcouple_output, lines)

    lines = add_end_to_namecouple()
    append_multiple_lines_to_file(namcouple_output, lines)

    # Creating the yaml_conf_name_ifs_toy_model 
    
    ifs_conf_yaml = toy_model_yaml_conf_output
    create_yaml_conf_name_ifs_toy_model(ifs_conf_yaml,ModelNameReceive, RunLengthSec, TimeStepSec,GridInfo,OasisOutputVars_readin)
