-- Copyright (c) 2015, Richard Laugesen (richard@tinyrock.com)
-- All rights reserved.

-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:

-- * Redistributions of source code must retain the above copyright notice, this
--   list of conditions and the following disclaimer.

-- * Redistributions in binary form must reproduce the above copyright notice,
--   this list of conditions and the following disclaimer in the documentation
--   and/or other materials provided with the distribution.

-- * Neither the name of Tiny Rock Pty Ltd nor the names of its
--   contributors may be used to endorse or promote products derived from
--   this software without specific prior written permission.

-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-- AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-- IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-- DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-- FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-- DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-- SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-- CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-- OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-- OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

module Main where

import qualified Data.ByteString.Lazy as BL
import qualified Data.Vector as V
import Control.Monad
import Control.Applicative
import Data.List
import Data.Csv
import Text.Printf
import System.Environment
import System.Random
    
import GR4J
import Stats
import Verification
import Optimisation

main :: IO ()
main = do
  [file] <- getArgs
  rows <- parseFile file
  randG <- newStdGen
         
  let obs = V.toList rows
      params = ParameterSet 400 0 150 5
      --params = ParameterSet 224.178126 0.874419 86.495168 1.114065
      --params = ParameterSet 873.747803784 0.48279895597 47.6515381972 1.05798738125
      init = initialState obs params

      sim = run obs params init

      qsim = [q | (q, _) <- sim]
      qobs = map q obs      
      states = [state | (_, state) <- sim]
      v1s = map v1 states
      v2s = map v2 states
      results = zip3 qsim v1s v2s

      optResults = take 1000 $ optimise obs params (kge qobs qsim) kge randG      
      (optParams, optObjFunValue) = last optResults

                   
  putStr "\nRunning GR4J simulation with forcings from test_data.csv\n"

  putStr $ printf "\nParameterSet: %s\n" (show params)
  putStr $ printf "Initial conditions: %s\n\n" (show init)
  putStr $ printf "First 10 time-steps of simulation (q, v1, v2):\n"
  putStr $ unlines $ take 10 $ map showResults results

  putStr $ printf "\nMean obs precipitation: %.3f mm\n" (mean (map p obs))
  putStr $ printf "Mean obs potential evapotranspiration: %.3f mm\n" (mean (map pe obs))
  putStr $ printf "Mean obs streamflow: %.3f mm\n"  (mean (map q obs))
  putStr $ printf "Median obs precipitation: %.3f mm\n" (median (map p obs))
  putStr $ printf "Median obs potential evapotranspiration: %.3f mm\n" (median (map pe obs))  
  putStr $ printf "Median obs streamflow: %.3f mm\n"  (median (map q obs))
  putStr $ printf "Mean sim streamflow: %.3f mm\n" (mean qsim)
  putStr $ printf "Median sim streamflow: %.3f mm\n" (median qsim)  
  putStr $ printf "Standard Deviation sim streamflow: %.3f mm\n" (std qsim)
  putStr $ printf "\nNSE: %.3f\n" (nse qobs qsim)
  putStr $ printf "KGE: %.3f\n" (kge qobs qsim)
  putStr $ printf "RMSE: %.3f mm\n" (rmse qobs qsim)
  putStr $ printf "MAE: %.3f mm\n" (mae qobs qsim)
  putStr $ printf "Bias: %.3f mm\n" (bias qobs qsim)
  putStr $ printf "Bias (Score): %.3f\n" (biasScore qobs qsim)                  
  putStr $ printf "Bias (Multiplicative): %.3f\n" (multBias qobs qsim)         
  putStr $ printf "Relative Variance: %.3f\n" (relativeVar qobs qsim)         
  putStr $ printf "Bias (MSE): %.3f mm\n" (biasMse qobs qsim)
  putStr $ printf "Volume Error: %.3f%%\n" (volumeError qobs qsim)         
  putStr $ printf "Correlation: %.3f\n" (cor qobs qsim)

  --putStr $ printf "\nObjective Function Trace:\n"
  --putStr $ unlines $ map (printf "%.3f") [objFunValue | (_, objFunValue) <- optResults]        

  putStr $ printf "\nOptimised Parameters: %s\n" (show optParams)
  putStr $ printf "Objective Function: %.3f\n" (optObjFunValue)
         
  putStr "\nFull time-series written to test_results.csv\n\n"
  BL.writeFile "test_results.csv" $ encode results

showResults :: (Double, Double, Double) -> String
showResults (a, b, c) = printf "%.5f, %.5f, %.5f" a b c

parseFile = liftM (either (error . show) (V.tail . V.init)) . parseCSVFromFile

parseCSVFromFile :: FilePath -> IO (Either String (V.Vector Observation))
parseCSVFromFile = fmap (decode HasHeader) . BL.readFile

instance FromRecord Observation where
   parseRecord v
          | V.length v == 5
          = Observation <$>
            v .!  0 <*>
            v .!  3 <*>
            v .!  4 <*>
            v .!  1
          | otherwise = mzero
