require(data.table)
require(raster)

generateShocksTable = function(dt, startYear, endYear, shockFrequencyMultiplier, makeMapFiles){
  shockTable = NULL
  cotterellRegions =fread(file.path(baseDataDir,"cotterellRegions.csv"), header=TRUE)
  noCountriesTab = data.table(table(cotterellRegions[,list(Region)]))[,list(Region=V1,nCountries = N)]
  
  ratesTab = merge(data.table(table(dt[,list(Year,Region,Sector)])),noCountriesTab, by=c('Region'))
  #Shock rate is number of shocks occurring in a region in a given year, shock rate scaled according to number of countries in a region
  ratesTab[,shockRate := N/nCountries]
  
  #probability of getting a shock rate of value x in year y
  shockProbs = data.table(table(ratesTab[N != 0,list(Region,Sector,shockRate)]))[,list(Region, Sector,shockRate, prob=N*shockFrequencyMultiplier/53)] #53 as that is number of years in historic time period
  shockProbs = rbind(shockProbs,shockProbs[,list(shockRate=0,prob=1-sum(prob)), by=c('Region','Sector')])[prob != 0]
  
  for(cropType in unique(dt$Sector)){
    #magnitudes of shocks take from global distribution per sector at a country level
    sectorHist = hist(dt[Sector==cropType,PropMagnitude],freq=FALSE)
    
    for(region in unique(dt$Region)){	
      #calculate years where shocks could happen using probability of shock occurring
      va = data.table(Year=seq(startYear,endYear,1),
                      shockRate=sample(shockProbs[Sector==cropType & Region==region,as.numeric(shockRate)],size=endYear-startYear+1, replace=TRUE, 
                                       prob=shockProbs[Sector==cropType & Region==region,prob]))[shockRate > 0]
      if(nrow(va)>0){
        #calculate for above years how many countries to shock 
        va[,numberOfCountriesToShock := sum(rbinom(noCountriesTab[Region==region,nCountries],size=1,prob=shockRate))]
        va=va[numberOfCountriesToShock>0]
        
        for(year in va[,unique(Year)]){
          #randomly choose countries within the region to shock and by what magnitude
          bins=with(sectorHist,sample(length(mids),va[Year==year,numberOfCountriesToShock],p=density,replace=TRUE))
          shockMagnitudes=runif(length(bins),sectorHist$breaks[bins],sectorHist$breaks[bins+1])
          countries = sample(cotterellRegions[Region==region, Country],va[Year==year,numberOfCountriesToShock])
          shockTable=rbind(shockTable, data.table(year,mapFilename=countries,crop=cropType,shockFactor=shockMagnitudes, region))	
        }
      }	
    }
  }
  
  if (makeMapFiles){
    makeMaps(shockTable, outputDir)
  }
  shockTable[,mapFilename := paste0(mapFilename,'.asc')]
  
  shockTable
}

##perhaps better to just keep a list of country maps rather than make them each time in this script
makeMaps=function(countriesTab, outputDir){
  
  country_codes = fread(file.path(baseDataDir, "country_codes4.csv"), header=TRUE)
  boundary = raster(file.path(baseDataDir, "halfdeg", "country_boundaries.asc"))
  crs(boundary) = '+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs'
  
  for(mapFilename in unique(countriesTab[,mapFilename])){
    b2 = boundary %in% country_codes[Country %in% mapFilename, numCode]
    values(b2)[values(b2) == 0] = NA
    fileN = filename=file.path(baseDataDir, "halfdeg","yieldshockmaps",paste0(mapFilename,'.asc'))
    if (!file.exists(fileN)){
      writeRaster(b2, filename=fileN, format='ascii', bylayer=TRUE, overwrite=TRUE)
    }
    
  }
}

#plot(raster("C:\\Users\\rhenry2\\eclipse-workspace\\plumv2\\data\\halfdeg\\yieldshockmaps\\United States of America.asc"))

plumDirectory = "/exports/csce/eddie/geos/groups/LURG/models/PLUM"
baseDataDir = "/exports/csce/eddie/geos/groups/LURG/models/PLUM/roslyns_build_area/plumv2/data"
baseOutputDir="/exports/csce/eddie/geos/groups/LURG/models/PLUM/output"

ensembleFile = commandArgs(trailingOnly = TRUE)[1]
scenarioTable = fread(file.path(plumDirectory,ensembleFile))
singleYieldShockFile = commandArgs(trailingOnly = TRUE)[2]
print(singleYieldShockFile)

if("SHOCKS_FREQUENCY_MULTIPLIER" %in% colnames(scenarioTable)){
  ensembleCounter = 0
  ensembleShockTable=NULL
  for(ensemble in unique(scenarioTable[,Ensemble])){
    
    shockFreqMultiplier = unique(scenarioTable[Ensemble==ensemble, SHOCKS_FREQUENCY_MULTIPLIER])
    
    if(length(shockFreqMultiplier) > 1){
      print("SHOCKS_FREQUENCY_MULTIPLIER is different within a single ensemble, using only the first values")
    }
    
    if(shockFreqMultiplier != 0){
      
      if ((singleYieldShockFile == 1 & ensembleCounter==0) | singleYieldShockFile != 1){
        
        ensembleShockTable=generateShocksTable(dt=fread(file.path(baseDataDir,"historicShocks.csv")), startYear=2010, endYear=2100, 
                                               shockFrequencyMultiplier=shockFreqMultiplier,makeMapFiles=TRUE)
      }
      
      for(scenario in scenarioTable[Ensemble==ensemble, Scenario]){
        
        ensDir=file.path(ensemble,scenario)
        
        if (dir.exists(file.path(baseOutputDir,ensDir))) {
          print(paste(ensDir, "exists.  Creating yield shock file"))
          write.csv(ensembleShockTable,file.path(baseOutputDir,ensDir,"yieldshocks.csv"), row.names=FALSE, quote=FALSE)
        } 
        else{
          print(paste(ensDir, "does not exist.  Stopping"))
        }
        
      }	
      ensembleCounter = ensembleCounter + 1
    }
  }
}else{
  print("No column for shock frequency multiplier in scenario table, not creating yieldShock files")
}