#include <iostream>
#include "config.h"
#include "shell.h"

#include "oasis-cpp-interface.h"
#include "OasisCoupler.h"
#include <stdio.h>
#include <stdlib.h>

using namespace std;

char OasisCoupler::modelName[]       = "lpjg";

// *** Received FROM IFS ***

// 1
char OasisCoupler::fieldTempName[]   = "T2M_LPJG";
int  OasisCoupler::fieldTempId       =  0;

// 2
char OasisCoupler::fieldPrecipName[] = "TPRE_LPJ";
int  OasisCoupler::fieldPrecipId     =  0;

// 3
char OasisCoupler::fieldSnowCName[] = "SNOC_LPJ";
int  OasisCoupler::fieldSnowCId     =  0;

// 4
char OasisCoupler::fieldSnowDName[] = "SNOD_LPJ";
int  OasisCoupler::fieldSnowDId     =  0;

// 5
char OasisCoupler::fieldST1LName[] = "ST1L_LPJ";
int  OasisCoupler::fieldST1LId     =  0;

// 6
char OasisCoupler::fieldST2LName[] = "ST2L_LPJ";
int  OasisCoupler::fieldST2LId     =  0;

// 7
char OasisCoupler::fieldST3LName[] = "ST3L_LPJ";
int  OasisCoupler::fieldST3LId     =  0;

// 8
char OasisCoupler::fieldST4LName[] = "ST4L_LPJ";
int  OasisCoupler::fieldST4LId     =  0;

// 9
char OasisCoupler::fieldSM1LName[] = "SM1L_LPJ";
int  OasisCoupler::fieldSM1LId     =  0;

// 10
char OasisCoupler::fieldSM2LName[] = "SM2L_LPJ";
int  OasisCoupler::fieldSM2LId     =  0;

// 11
char OasisCoupler::fieldSM3LName[] = "SM3L_LPJ";
int  OasisCoupler::fieldSM3LId     =  0;

// 12
char OasisCoupler::fieldSM4LName[] = "SM4L_LPJ";
int  OasisCoupler::fieldSM4LId     =  0;

// 13
char OasisCoupler::fieldSWRadName[] = "SWNR_LPJ";
int  OasisCoupler::fieldSWRadId = 0;

// 14
char OasisCoupler::fieldLWRadName[] = "LWNR_LPJ";
int  OasisCoupler::fieldLWRadId = 0;


// *** Send TO IFS ***

// 16 
char OasisCoupler::fieldLowlaiName[] = "GUE_LLAI";
int  OasisCoupler::fieldLowlaiId     =  0;

// 17
char OasisCoupler::fieldHighlaiName[] = "GUE_HLAI";
int  OasisCoupler::fieldHighlaiId     =  0;

// 18
char OasisCoupler::fieldTypeHName[] = "GUE_TYPH";
int  OasisCoupler::fieldTypeHId     =  0;

// 19
char OasisCoupler::fieldFracHName[] = "GUE_FRAH";
int  OasisCoupler::fieldFracHId     =  0;

// 20
char OasisCoupler::fieldTypeLName[] = "GUE_TYPL";
int  OasisCoupler::fieldTypeLId     =  0;

// 21
char OasisCoupler::fieldFracLName[] = "GUE_FRAL";
int  OasisCoupler::fieldFracLId     =  0;

// *** Received FROM TM5 ***

// 22
char OasisCoupler::fieldCO2Name[] = "CO2_LPJG";
int  OasisCoupler::fieldCO2Id     =  0;

// *** Sent TO TM5 ***

// 23
char OasisCoupler::fieldCfluxNatName[] = "GUE_CNAT";
int  OasisCoupler::fieldCfluxNatId     =  0;

// 24
char OasisCoupler::fieldCfluxAntName[] = "GUE_CANT";
int  OasisCoupler::fieldCfluxAntId     =  0;

// 25
char OasisCoupler::fieldNPPName[] = "GUE_CNPP";
int  OasisCoupler::fieldNPPId     =  0;

// ecev3 - returns localcomm
int OasisCoupler::init(int &localcomm)
{
    int ierror;

    int componentId = -99;

    dprintf("OasisCoupler::init init_comp \n");

    ierror = OASISMCT::init_comp(&componentId,modelName);
    if (ierror != OASISMCT::OASIS_Ok) {
	dprintf("OasisCoupler::init init_comp ERROR!!\n");
	return ierror;
	}

	// ecev3 - for internal parallelisation
	int localCommId = -99;
    	ierror = OASISMCT::get_localcomm(localCommId);
    	if (ierror != OASISMCT::OASIS_Ok) {
		dprintf("OasisCoupler::get_localcomm ERROR - %i\n", ierror);
		OASISMCT::terminate();
		return ierror;
	}

	// Save the local comm index
	localcomm = localCommId;

	return OASISMCT::OASIS_Ok;
};


// ecev3 - returns couplcomm
int OasisCoupler::create_couplcomm(int myrank, int lcomm, int& ccomm) {

	int ierror;

	dprintf("OasisCoupler::create_couplcomm \n");

	int cplcomm = -99;
	ierror = OASISMCT::create_couplcomm(myrank, lcomm, &cplcomm);
	if ( ierror != OASISMCT::OASIS_Ok ) 
		return ierror;
	else 
		ccomm = cplcomm;

    return OASISMCT::OASIS_Ok;
};



// ecev3 
int OasisCoupler::init_part_defvar(int nx, int ny, bool TM5coupled, int rank) {

	int ierror;

    int partitionId      = -99;
    int partitionDescr[] = {0,0,nx*ny};

	// ecev3 - see OASIS-MCT documentation
	if (rank != 0)
		partitionDescr[2] = 0;

	dprintf("OasisCoupler::init_part_defvar def_partition for rank %i\n", rank);

    ierror = OASISMCT::def_partition(&partitionId,partitionDescr,3);
	//  Bug in OASISMCT: def_partition doesn't return a valid return value
	if ( ierror != OASISMCT::OASIS_Ok ) { 
		dprintf("OasisCoupler::init_part_defvar: def_partition returned %i\n", ierror);
		return ierror;
	}

	dprintf("OasisCoupler::init_part_defvar: rank %i partitionId is %i\n", rank, partitionId);

    int varNumDims[] = {2,1};
    int varShape[]   = {1,nx,1,ny};


	dprintf("OasisCoupler::init_part_defvar def_var statements for rank %i \n", rank);

	// *** Received FROM IFS (Currently 14) ***

	ierror = OASISMCT::def_var(&fieldTempId,fieldTempName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldPrecipId,fieldPrecipName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldSnowCId,fieldSnowCName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldSnowDId,fieldSnowDName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldST1LId,fieldST1LName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldST2LId,fieldST2LName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldST3LId,fieldST3LName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldST4LId,fieldST4LName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldSM1LId,fieldSM1LName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldSM2LId,fieldSM2LName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldSM3LId,fieldSM3LName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldSM4LId,fieldSM4LName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldSWRadId,fieldSWRadName,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldLWRadId,fieldLWRadName,partitionId,varNumDims,OASISMCT::OASIS_In, varShape, OASISMCT::OASIS_Real);
	if (ierror != OASISMCT::OASIS_Ok) return ierror;

	// *** Sent TO IFS (Currently 6) ***

	ierror = OASISMCT::def_var(&fieldLowlaiId,fieldLowlaiName,partitionId,varNumDims,OASISMCT::OASIS_Out,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldHighlaiId,fieldHighlaiName,partitionId,varNumDims,OASISMCT::OASIS_Out,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldTypeHId,fieldTypeHName,partitionId,varNumDims,OASISMCT::OASIS_Out,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldFracHId,fieldFracHName,partitionId,varNumDims,OASISMCT::OASIS_Out,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldTypeLId,fieldTypeLName,partitionId,varNumDims,OASISMCT::OASIS_Out,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	ierror = OASISMCT::def_var(&fieldFracLId,fieldFracLName,partitionId,varNumDims,OASISMCT::OASIS_Out,varShape,OASISMCT::OASIS_Real);
	if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

	// *** TM5 Coupling (Currently 4 fields) ***

	if (TM5coupled) {

		dprintf("OasisCoupler::init_part_defvar - initialising TM5 coupling\n");

		// FROM TM5
		ierror = OASISMCT::def_var(&fieldCO2Id,fieldCO2Name,partitionId,varNumDims,OASISMCT::OASIS_In,varShape,OASISMCT::OASIS_Real);
		if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

		// TO TM5
		ierror = OASISMCT::def_var(&fieldCfluxNatId,fieldCfluxNatName,partitionId,varNumDims,OASISMCT::OASIS_Out,varShape,OASISMCT::OASIS_Real);
		if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

		// TO TM5
		ierror = OASISMCT::def_var(&fieldCfluxAntId,fieldCfluxAntName,partitionId,varNumDims,OASISMCT::OASIS_Out,varShape,OASISMCT::OASIS_Real);
		if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

		// TO TM5
		ierror = OASISMCT::def_var(&fieldNPPId,fieldNPPName,partitionId,varNumDims,OASISMCT::OASIS_Out,varShape,OASISMCT::OASIS_Real);
		if ( ierror != OASISMCT::OASIS_Ok ) return ierror;
	}

	dprintf("OasisCoupler::init_part_defvar - calling OASISMCT::enddef on rank %i\n",rank);
	ierror = OASISMCT::enddef();
	dprintf("OasisCoupler::init_part_defvar - OASISMCT::enddef returned ierror %i on rank %i\n",ierror, rank);

    if (ierror != OASISMCT::OASIS_Ok) return ierror;

    return OASISMCT::OASIS_Ok;
};


int OasisCoupler::abort(int compid, std::string routine_name, std::string abort_message, int return_code)
{
	//dprintf("OasisCoupler::abort %s\n",(char)abort_message);
    int ierror = OASISMCT::abort(compid, routine_name, abort_message, return_code);
    if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

    return 0;
};

int OasisCoupler::finalize(void)
{
    dprintf("OasisCoupler::finalize\n");
    int ierror = OASISMCT::terminate();
    if ( ierror != OASISMCT::OASIS_Ok ) return ierror;

    return 0;
};



// Up to 18 fields (17 from IFS + 1 from TM5 if TM5coupled == true)
int OasisCoupler::couple_get(int time, int nx, int ny, bool TM5coupled, double * fieldTempData, double * fieldPrecipData, 
	double * fieldSnowCData, double * fieldSnowDData,
	double * fieldST1LData, double * fieldST2LData, double * fieldST3LData, double * fieldST4LData, 
	double * fieldSM1LData, double * fieldSM2LData, double * fieldSM3LData, double * fieldSM4LData, 
	double * fieldSWRadData, double * fieldLWRadData, double * fieldCO2Data) {

    int ierror = 0;

	try {

		// printf("  OasisCoupler::couple_get -- for time %i\n",time);
		ierror = OASISMCT::get_2d(OasisCoupler::fieldTempId,time,fieldTempData,nx,ny);
		// printf("  OasisCoupler::couple_get -- got T2m\n");
	}
	catch(...) {

		switch (ierror)
		{
			case OASISMCT::OASIS_Recvd:
 				dprintf(" T2M received from other model: %i \n", ierror);
 				break;
			case OASISMCT::OASIS_FromRest:
 				dprintf(" Read from restart file only: %i \n", ierror);
 				break;
			case OASISMCT::OASIS_Input:
 				dprintf(" Read from input file only: %i \n", ierror);
 				break;
			case OASISMCT::OASIS_RecvOut:
 				dprintf(" Received from other model and written to an output file: %i \n", ierror);
 				break;
			case OASISMCT::OASIS_FromRestOut:
 				dprintf(" Read from restart file and written to output file: %i \n", ierror);
 				break;
			case OASISMCT::OASIS_Ok:
 				dprintf(" OASIS_Ok, no error. T2M received OK, continue: %i \n", ierror);
 				break;
			default:
 				dprintf(" OasisCoupler::couple_get - unknown ierror reported from OASIS_GET: %i \n", ierror);
				break;
		}
    
	}


	ierror = OASISMCT::get_2d(OasisCoupler::fieldSWRadId, time, fieldSWRadData, nx, ny);
	//printf("  OasisCoupler::couple -- got Rad\n");

	ierror = OASISMCT::get_2d(OasisCoupler::fieldLWRadId, time, fieldLWRadData, nx, ny);
	//printf("  OasisCoupler::couple -- got Rad\n");

    ierror = OASISMCT::get_2d(OasisCoupler::fieldSM1LId,time,fieldSM1LData,nx,ny);
    // printf("  OasisCoupler::couple -- got SM1\n");
    ierror = OASISMCT::get_2d(OasisCoupler::fieldSM2LId,time,fieldSM2LData,nx,ny);
    // printf("  OasisCoupler::couple -- got SM2\n");
    ierror = OASISMCT::get_2d(OasisCoupler::fieldSM3LId,time,fieldSM3LData,nx,ny);
    // printf("  OasisCoupler::couple -- got SM3\n");
    ierror = OASISMCT::get_2d(OasisCoupler::fieldSM4LId,time,fieldSM4LData,nx,ny);
    // printf("  OasisCoupler::couple -- got SM4\n");
    
    ierror = OASISMCT::get_2d(OasisCoupler::fieldST1LId,time,fieldST1LData,nx,ny);
    // printf("  OasisCoupler::couple -- got ST1\n");
    ierror = OASISMCT::get_2d(OasisCoupler::fieldST2LId,time,fieldST2LData,nx,ny);
    // printf("  OasisCoupler::couple -- got ST2\n");
    ierror = OASISMCT::get_2d(OasisCoupler::fieldST3LId,time,fieldST3LData,nx,ny);
    // printf("  OasisCoupler::couple -- got ST3\n");
    ierror = OASISMCT::get_2d(OasisCoupler::fieldST4LId,time,fieldST4LData,nx,ny);
    // printf("  OasisCoupler::couple -- got ST4\n");

    ierror = OASISMCT::get_2d(OasisCoupler::fieldSnowCId,time,fieldSnowCData,nx,ny);
    //printf("  OasisCoupler::couple -- got SnowC\n");
    ierror = OASISMCT::get_2d(OasisCoupler::fieldSnowDId,time,fieldSnowDData,nx,ny);
    // printf("  OasisCoupler::couple -- got SnowD\n");

    ierror = OASISMCT::get_2d(OasisCoupler::fieldPrecipId,time,fieldPrecipData,nx,ny);
    // printf("  OasisCoupler::couple -- got Prec\n");

	// ecev3 - only receive CO2 fields from TM5 if we're coupled 
	if (TM5coupled) {
		ierror = OASISMCT::get_2d(OasisCoupler::fieldCO2Id,time,fieldCO2Data,nx,ny);
		dprintf(" OasisCoupler::couple -- got CO2 -- cell 5000 value is %f\n",fieldCO2Data[5000]);
	}

    return 0;
};



// Up to 8 fields (6 to IFS + 2 to TM5 if TM5coupled == true)
int OasisCoupler::couple_put(int time, int nx, int ny, bool TM5coupled, double * fieldLowlaiData, double * fieldHighlaiData,
							 double * fieldLPJGtypeHData, double * fieldLPJGfracHData, 
							 double * fieldLPJGtypeLData, double * fieldLPJGfracLData,
							 double * fieldLPJGCfluxNatData, double * fieldLPJGCfluxAntData,
							 double * fieldLPJGNPPData) {

    int ierror = 0;

	try {
	    // printf("  OasisCoupler::couple_put -- for time %i\n",time);
		ierror = OASISMCT::put_2d(OasisCoupler::fieldLowlaiId,time,fieldLowlaiData,nx,ny);
	}
	catch (...) {

		switch (ierror)
		{
			case OASISMCT::OASIS_Sent:
 				dprintf(" LLAI field sent to the other model: %i \n",ierror);
 				break;
			case OASISMCT::OASIS_LocTrans:
 				dprintf(" OASIS_LocTrans: %i \n",ierror);
 				break;
			case OASISMCT::OASIS_Output:
 				dprintf(" Output file only: %i \n",ierror);
 				break;
			case OASISMCT::OASIS_ToRest:
 				dprintf(" Written to restart file only: %i \n",ierror); 
 				break;
			case OASISMCT::OASIS_ToRestOut:
 				dprintf(" Written to restart file and written to output file: %i \n",ierror);
 				break;
			case OASISMCT::OASIS_SentOut:
				dprintf(" OASIS_SentOut: %i \n",ierror);
				break;
			case OASISMCT::OASIS_Ok:
 				dprintf(" OASIS_Ok, LLAI sent, no error: %i - continue\n",ierror);
 				break;
			default:
 				dprintf(" OasisCoupler::couple_put - unknown ierror %i reported from OASIS_PUT for field %i \n",ierror,OasisCoupler::fieldLowlaiId);
				break;
		}
	}

	// ecev3 - always send fields to IFS
	ierror = OASISMCT::put_2d(OasisCoupler::fieldHighlaiId,time,fieldHighlaiData,nx,ny);
	if (ierror != OASISMCT::OASIS_Ok)
		dprintf(" OasisCoupler::couple_put HLAI - ierror: %i\n",ierror);

	ierror = OASISMCT::put_2d(OasisCoupler::fieldTypeHId,time,fieldLPJGtypeHData,nx,ny);
	if (ierror != OASISMCT::OASIS_Ok)
		dprintf(" OasisCoupler::couple_put TYPH - ierror: %i\n",ierror);

	ierror = OASISMCT::put_2d(OasisCoupler::fieldFracHId,time,fieldLPJGfracHData,nx,ny);
	if (ierror != OASISMCT::OASIS_Ok)
		dprintf(" OasisCoupler::couple_put FRACH - ierror: %i\n",ierror);

	ierror = OASISMCT::put_2d(OasisCoupler::fieldTypeLId,time,fieldLPJGtypeLData,nx,ny);
	if (ierror != OASISMCT::OASIS_Ok)
		dprintf(" OasisCoupler::couple_put TYPL - ierror: %i\n",ierror);

	ierror = OASISMCT::put_2d(OasisCoupler::fieldFracLId,time,fieldLPJGfracLData,nx,ny);
	if (ierror != OASISMCT::OASIS_Ok)
		dprintf(" OasisCoupler::couple_put FRACL - ierror: %i\n",ierror);

	// ecev3 - only send CO2 fluxes to TM5 if we are coupled 
	if (TM5coupled) {
		ierror = OASISMCT::put_2d(OasisCoupler::fieldCfluxNatId,time,fieldLPJGCfluxNatData,nx,ny);
		if (ierror != OASISMCT::OASIS_Ok)
			dprintf(" OasisCoupler::couple_put CFLUXNAT - ierror: %i\n",ierror);
	 
		ierror = OASISMCT::put_2d(OasisCoupler::fieldCfluxAntId,time,fieldLPJGCfluxAntData,nx,ny);
		if (ierror != OASISMCT::OASIS_Ok)
			dprintf(" OasisCoupler::couple_put CFLUXANT - ierror: %i\n",ierror);
	 
		ierror = OASISMCT::put_2d(OasisCoupler::fieldNPPId,time,fieldLPJGNPPData,nx,ny);
		if (ierror != OASISMCT::OASIS_Ok)
			dprintf(" OasisCoupler::couple_put NPP - ierror: %i\n",ierror);
	}

    return 0;
};