{-# OPTIONS_GHC -Wall #-}

{-|

  Module      : AutoBench.Internal.Analysis
  Description : Coordinates the analysis of runtime measurements.
  Copyright   : (c) 2018 Martin Handley
  License     : BSD-style
  Maintainer  : martin.handley@nottingham.ac.uk
  Stability   : Experimental
  Portability : GHC

  This module is responsible for coordinating the statistical analysis of
  runtime measurements. Raw measurements are subject to regression analysis
  (ridge regression) and linear models are fitted sets of (input size, runtime)
  coordinates. Regression models are then ranked according to fitting statistics
  generated by e.g., cross-validation (see 'Stats'). The equations of best
  fitting models are used to approximate the time complexity of test programs.

  * 'quickAnalysis' is for analysing the results of QuickBenching;
  * 'analyse' and 'analyseWith' are for analysing the results of AutoBenching.

  Note that 'analyse' and 'analyseWith' will be moved to a user exposed module
  when 'TestReport's can be saved to/loaded from file.

-}

{-
   ----------------------------------------------------------------------------
   <TO-DO>:
   ----------------------------------------------------------------------------
   - Address the cleansing of input sizes to remove 0;
   -
-}

module AutoBench.Internal.Analysis
  (

  -- * QuickBench.
    quickAnalyseWith   -- Analyse the given @[QuickReport]@, i.e., test results using the given 'AnalOpts'.
  -- * AutoBench.
  , analyse            -- Analyse the given 'TestReport', i.e., test results using the default 'AnalOpts'.
  , analyseWith        -- Analyse the given 'TestReport', i.e., test results using the given 'AnalOpts'

  ) where

import           Control.Arrow         ((&&&))
import           Criterion.Types       (OutlierEffect(..))
import           Data.Default          (def)
import           Data.List             (genericLength, sort, sortBy)
import           Data.Maybe            (catMaybes)
import qualified Data.Vector.Storable  as V
import           Numeric.LinearAlgebra (Vector, norm_1, norm_2)
import           System.IO.Unsafe      (unsafePerformIO)
import           System.Random         (randomRIO)

import AutoBench.Internal.AbstractSyntax  (Id)
import AutoBench.Internal.Configuration   (maximumModelPredictors)
import AutoBench.Internal.UserIO          (outputAnalysisReport, outputQuickAnalysis)
import AutoBench.Internal.Regression      (generateLinearCandidate, fitRidgeRegress)
import AutoBench.Internal.UserInputChecks (validateAnalOpts, validateTestReport)
import AutoBench.Internal.Utils           (notNull, uniqPairs)

import AutoBench.Internal.Types
  ( AnalOpts(..)
  , AnalysisReport(..)
  , BenchReport(..)
  , Coord
  , Coord3
  , CVStats(..)
  , Improvement
  , InputError(..)
  , LinearCandidate(..)
  , LinearFit(..)
  , SimpleReport(..)
  , SimpleResults(..)
  , Stats(..)
  , TestReport(..)
  , QuickAnalysis(..)
  , QuickReport(..)
  , QuickResults(..)
  , numPredictors
  , simpleReportsToCoords
  )

-- * QuickBench Top-level

-- | Analyse the given @[QuickReport]@, i.e., test results using the
-- given 'AnalOpts'. Analysis results are output to the console and file
-- depending on the 'AnalOpts'. The boolean parameter is whether the
-- test programs have the same results according to QuickCheck testing.
-- Note: this function is not exposed to users.
quickAnalyseWith :: AnalOpts -> Bool -> [QuickReport] -> IO ()
quickAnalyseWith _ _ [] = print $ QuickBenchErr "No results to analyse."
quickAnalyseWith aOpts eql qrs = outputQuickAnalysis aOpts eql $
  quickAnalysis aOpts qrs

-- * AutoBench Top-level

-- | Analyse the given 'TestReport', i.e., test results using the
-- default 'AnalOpts'.
analyse :: TestReport -> IO ()
analyse  = analyseWith def

-- | Analyse the given 'TestReport', i.e., test results using the
-- given 'AnalOpts'. Analysis results are output to the console and file
-- depending on the 'AnalOpts'.
analyseWith :: AnalOpts -> TestReport -> IO ()
analyseWith aOpts tr
  | notNull (aOptsErrs ++ trErrs) = do
      putStrLn "Cannot analyse results due to one or more errors:"
      mapM_ print (aOptsErrs ++ trErrs)
  -- Just output the results of statistical analysis.
  | otherwise = outputAnalysisReport aOpts tr $ calculateAnalysisReport aOpts tr

  where
    -- Validate the 'AnalOpts'.
    aOptsErrs = validateAnalOpts aOpts
    -- Validate the 'TestReport'
    trErrs    = validateTestReport tr

-- * QuickBench Analysis

-- | Perform statistical analysis on the benchmarking results (i.e., runtime
-- measurements of test programs) in the given @[QuickReport]@ and produce a
-- 'QuickAnalysis' report to summarise the system's analysis phase.
quickAnalysis :: AnalOpts -> [QuickReport] -> QuickAnalysis
quickAnalysis aOpts qrs =
  QuickAnalysis
    {
      _qAnlys = quickResults aOpts qrs                               -- A set of quick results ('QuickResults') for each test program.
    , _qImps  = catMaybes $ fmap (uncurry $ calculateImprovements    -- A set of improvements.
        (_improv aOpts)) (uniqPairs $ zip names coords)
    }
  where
    -- Names of test programs
    names = fmap _qName qrs
    -- (Input size(s), runtime) coordinates.
    coords = fmap _qRuntimes qrs

-- | Perform statistical analysis on the benchmarking results (i.e., runtime
-- measurements of test programs) in the given @[QuickReport]@ and produce an
-- analysis report ('QuickResults') /for each test program/.
quickResults :: AnalOpts -> [QuickReport] -> [QuickResults]
quickResults aOpts = fmap quickResult
  where
    quickResult :: QuickReport -> QuickResults
    quickResult qr =
      QuickResults
        {
          _qrIdt  = _qName qr
        -- Raw measurements.
        , _qrRaws = coords
        -- Fit all models in the 'AnalOpts' to the raw measurements.
        , _qrFits = fitCoords aOpts coords
        }
      where coords = _qRuntimes qr

-- * AutoBench Analysis

-- | Perform statistical analysis on the benchmarking results (i.e., runtime
-- measurements of test programs) in the given 'TestReport' and produce an
-- 'AnalysisReport' to summarise the system's overall analysis phase.
calculateAnalysisReport :: AnalOpts -> TestReport -> AnalysisReport
calculateAnalysisReport aOpts tr =
  AnalysisReport
    {
      _anlys = anlys                                               -- A set of simple results ('SimpleResults') for each test program.
    , _imps  = catMaybes $ fmap (uncurry $ calculateImprovements   -- A set of improvements.
        -- Compare every set of measurements against every other set.
        (_improv aOpts)) (uniqPairs $ zip srNames srCoords)
    , _blAn  = blAn                                                -- Analysis of baseline measurements, if applicable.
    }
  where
    -- 'SimpleResults' for test programs and baseline measurements, if applicable.
    (anlys, blAn) =  calculateSimpleResults aOpts tr

    -- For generating improvements: compare all results from test programs
    -- against eachother using 'uniqPairs'.

    -- SimpleReports.
    srs = _reports (_br tr)
    -- Names of test programs from 'SimpleReport's.
    srNames  = fmap (_name . head) srs
    -- (Input size(s), runtime) coordinates.
    srCoords = zipWith simpleReportsToCoords srNames srs

-- | Perform statistical analysis on the benchmarking results (i.e., runtime
-- measurements of test programs) in the given 'TestReport' and produce an
-- analysis report ('SimpleResults') /for each test program/.
-- If baseline measurements have been taken then also produce a 'SimpleReport'
-- for these results.
calculateSimpleResults
  :: AnalOpts
  -> TestReport
  -> ([SimpleResults], Maybe SimpleResults) -- Maybe is for baseline measurements.
calculateSimpleResults aOpts tr = (progResults, blResults)

  where
    -- Analysis for test programs.
    progResults = zipWith calculateSimpleResult (_tProgs tr) (_reports $ _br tr)
    -- Analysis for baseline measurements.
    blResults = case _baselines $ _br tr of
      []  -> Nothing
      bls -> Just $ calculateSimpleResult "Baseline measurements" bls                               -- <TO-DO>: This shouldn't be a string.

    -- For a set of 'SimpleReport's (i.e., test cases) relating to a
    -- given test program, perform statistical analysis and generate
    -- 'SimpleResults'.
    calculateSimpleResult
      :: Id               -- Name of test program.
      -> [SimpleReport]   -- Benchmarking measurements for each test case.
      -> SimpleResults
    calculateSimpleResult idt srs =
      SimpleResults
        {
          _srIdt           = idt
        -- Raw measurements.
        , _srRaws          = coords
        -- Cumulative statistics for /all/ test cases, taken from Criterion's 'Report'.
        , _srStdDev        = stdDev
        , _srAvgOutVarEff  = avgOutVarEff
        , _srAvgPutVarFrac = avgOutVarFrac
        -- Fit all models in the 'AnalOpts' to the raw measurements.
        , _srFits          = fitCoords aOpts coords
        }
      where
        coords = simpleReportsToCoords idt srs
        (stdDev, avgOutVarFrac, avgOutVarEff) = simpleReportSummary srs

    -- Provide an summary of the simple statistics taken from
    -- Criterion's 'Report's relating to the /same/ test program:
    --
    -- * Standard deviation of all test cases;
    -- * Average effect of outliers on variance (percentage);
    -- * Average effect of outliers on variance ('OutlierEffect').
    simpleReportSummary
      :: [SimpleReport]
      -> (Double, Double, OutlierEffect)
    simpleReportSummary []  = (0, 0, Unaffected)                                                    -- <TO-DO> What is this error handling?
    simpleReportSummary srs = (stdDev srs, avgOutVarFrac, avgOutVarEff)
      where
        -- Standard deviation for all test cases: sqrt (SUM variance_i/samples_i).
        stdDev = sqrt . sum . fmap (uncurry (/) . ((** 2) .
          _stdDev &&& fromIntegral . _samples))

        -- Average effect of outliers on variance.
        avgOutVarFrac = sum (fmap _outVarFrac srs) / genericLength srs
        avgOutVarEff
          | avgOutVarFrac < 0.01 = Unaffected -- Same as Criterion.
          | avgOutVarFrac < 0.1  = Slight
          | avgOutVarFrac < 0.5  = Moderate
          | otherwise            = Severe

-- * Linear fitting

-- | Fit a 'LinearCandidate' model to a data set, generating a 'LinearFit'
-- that can be used to asses the model's goodness of fit.
candidateFit
  :: ([Coord] -> LinearCandidate -> Vector Double)  -- Fitting function.
  -> Double                                         -- Cross-validation train/validate data split.
  -> Int                                            -- Cross-validation iterations.
  -> [Coord]                                        -- Data set.
  -> LinearCandidate                                -- Model to fit.
  -> Maybe LinearFit
candidateFit ff split iters coords c
  | length cleansedCoords < maximumModelPredictors + 1 = Nothing     -- Need a minimum of (maximumModelPredictors + 1) coordinates.
  | V.null coeffs = Nothing                                          -- If coefficients are null, can't use model.
  | otherwise = Just
      LinearFit
        {
          _lft  = _lct c
        , _cfs  = coeffs
        , _ex   = _fex   c coeffs
        , _yhat = _fyhat c coeffs
        , _sts  = sts
        }
  where
    -- Model coefficients when fit to /entire/ data set.
    -- This is for e.g., R^2, Adj. R^2, BIC, AIC CP.
    coeffs = ff cleansedCoords c                                                                    -- <TO-DO>: can 'ff' fail?
    -- Statistics from cross-validation, e.g., PRESS, PMAE, PMSE.
    cvSts  = cvCandidateFit ff split iters cleansedCoords c
    -- Combine statistics from each cross-validation iteration and also compute
    -- other goodness of fit measures.
    sts    = stats c cleansedCoords coeffs cvSts

    -- Cleanse in case transformation produces infinite values.                                     -- <TO-DO> *** This needs addressing. ***
    -- /This happens when for size 0 only, LogBase _ 0 = -Infinity/.
    -- Remove the @(0, _)@ coordinate.
    -- ** Only cleanse runtimes predicted by models. **
    cleansedCoords = cleanse coords
    cleanse :: [Coord] -> [Coord]
    cleanse  = filter (\(x, _) -> x /= 0)

-- ** Linear fitting cross-validation

-- | Fit a 'LinearCandidate' to a data set using Monte Carlo cross-validation.
-- Return the fitting statistics ('CVStats') from each iteration of
-- cross-validation.
cvCandidateFit
  :: ([Coord] -> LinearCandidate -> Vector Double)  -- Fitting function (ridge regression in practice).
  -> Double                                         -- Train/evaluate data split.
  -> Int                                            -- Number of iterations.
  -> [Coord]                                        -- Full data set.
  -> LinearCandidate                                -- Model to fit.
  -> [CVStats]
cvCandidateFit ff split iters coords c =
  fmap (uncurry trainAndValidate) splits
  where
    splits = fmap (splitRand split) (replicate iters coords)

    -- Train the model on the training data and evaluate it on the validation
    -- data: calculating the cross-validation statistics ('CVStats').
    trainAndValidate :: [Coord] -> [Coord] -> CVStats
    trainAndValidate train valid = cvStats c valid (ff train c)

-- ** Linear fitting statistics

-- | Fitting statistics generated from a /single/ iteration of cross-validation.
-- See 'stats' for cumulative fitting statistics.
cvStats
  :: LinearCandidate    -- The model.
  -> [Coord]            -- The /evaluation/ set of this iteration of cross-validation.
  -> Vector Double      -- The coefficients of the model calculated using the /training/ set.
  -> CVStats            -- The model's fitting statistics for this iteration of cross-validation.
cvStats c coords coeffs =
  CVStats
   {
     _cv_mse    = mse
   , _cv_mae    = mae
   , _cv_ss_tot = ss_tot
   , _cv_ss_res = ss_res
   }
  where
    ---------------------------------------------------------------------------
    (xs, ys) = unzip coords
    _X       = V.fromList xs                              -- X. ** don't transform using fxs **
    _Y       = V.fromList ys                              -- Y.
    n        = length coords
    n'       = fromIntegral n                             -- n data points.
    ---------------------------------------------------------------------------
    yhat    = _fyhat c coeffs _X                          -- Model y-coords.
    ybar    = V.sum _Y / n'                               -- Average of data y-coords.
    ybars   = V.replicate n ybar
    ---------------------------------------------------------------------------
    mse     = ss_res / n'                                 -- Mean squared error       ss_res / n
    mae     = (norm_1 $ V.zipWith (-) _Y yhat) / n'       -- Mean absolute error:     (SUM |y_i - f(y_i)|) / n
    ss_tot  = (norm_2 $ V.zipWith (-) _Y ybars) ** 2.0    -- Total sum of squares:    SUM (y_i - avg(y))^2
    ss_res  = (norm_2 $ V.zipWith (-) _Y yhat)  ** 2.0    -- Residual sum of squares: SUM (y_i - f(y_i))^2

-- | Overall fitting statics for a model and a given data set. Some statistics
-- are cumulative fitting stats, generated by cross-validation (PMSE, PMAE,
-- PRESS, SST, Pred. R^2); others are generated by using the whole data set as
-- training data (R^2, Adj. R^2, AIC, BIC, Mallow's CP).
stats
  :: LinearCandidate  -- The model.
  -> [Coord]          -- The data set
  -> Vector Double    -- The coefficients of the model calculated using the /entire/ set.
  -> [CVStats]        -- The statistics from each iteration of cross-validation.
  -> Stats            -- The cumulative fitting statistics.
stats c coords coeffs cvSts = _sts
  where
    -- Statistics calculated below.
    _sts = Stats
            {
              _p_mse    = p_mse
            , _p_mae    = p_mae
            , _ss_tot   = cv_ss_tot
            , _p_ss_res = p_ss_res
            , _r2       = r2
            , _a_r2     = a_r2
            , _p_r2     = p_r2
            , _bic      = bic
            , _aic      = aic
            , _cp       = cp
            }
    ---------------------------------------------------------------------------
    -- Statistics for /complete/ data set:
    ---------------------------------------------------------------------------
    (xs, ys) = unzip coords
    _X       = V.fromList xs                                           -- X. ** do not transform using fxs **
    _Y       = V.fromList ys                                           -- Y.
    n        = length coords
    n'       = fromIntegral n                                          -- n data points.
    k'       = fromIntegral (numPredictors $ _lct c)                   -- k predictors.
    ---------------------------------------------------------------------------
    yhat     = _fyhat c coeffs _X                                      -- Model y-coords.
    ybar     = V.sum _Y / n'                                           -- Average of data y-coords.
    ybars    = V.replicate n ybar
    ---------------------------------------------------------------------------
    _mse     = _ss_res / n'                                            -- Mean squared error:                       ss_res / n
    _ss_tot  = (norm_2 $ V.zipWith (-) _Y ybars) ** 2.0                -- Total sum of squares:                     SUM (y_i - avg(y))^2
    _ss_res  = (norm_2 $ V.zipWith (-) _Y yhat)  ** 2.0                -- Residual sum of squares:                  SUM (y_i - f(y_i))^2
    r2       = 1.0 - (_ss_res / _ss_tot)                               -- Coefficient of Determination:             1 - ss_res/ss_tot
    a_r2     = 1.0 - ((1.0 - r2) * (n' - 1.0) / (n' - k' - 1.0))       -- Adjusted Coefficient of Determination:    1 - ((1 - r^2)(n - 1)/(n - k - 1))
    ---------------------------------------------------------------------------
    -- http://www.stat.wisc.edu/courses/st333-larget/aic.pdf (implementations from R).
    ---------------------------------------------------------------------------
    aic = (n' * log (_ss_res / n')) + (2 * k')                         -- Akaike’s Information Criterion:           n * ln(ss_res/n) + 2k
    bic = (n' * log (_ss_res / n')) + (k' * log n')                    -- Bayesian Information Criterion:           n * ln(ss_res/n) + k * ln(n)
    cp  = (_ss_res / _mse) - (n' - (2 * k'))                           -- Mallows' Cp Statistic:                    (ss_res/_mse) - (n - 2k)
    ---------------------------------------------------------------------------
    -- Cross-validated statistics:
    ---------------------------------------------------------------------------
    cv_n      = genericLength cvSts
    p_mse     = (sum $ fmap _cv_mse cvSts)    / cv_n                   -- 1/n * SUM mse_k         for each iteration k
    p_mae     = (sum $ fmap _cv_mae cvSts)    / cv_n                   -- 1/n * SUM mae_k         for each iteration k
    cv_ss_tot = (sum $ fmap _cv_ss_tot cvSts) / cv_n                   -- 1/n * SUM ss_tot_k      for each iteration k
    p_ss_res  = (sum $ fmap _cv_ss_res cvSts) / cv_n                   -- 1/n * SUM ss_res_k      for each iteration k
    p_r2     = 1.0 - (p_ss_res / cv_ss_tot)                            -- 1 - p_ss_res/cv_ss_tot

-- * Helpers

-- | Split a list into two random sublists of specified sizes.
-- Warning: uses 'unsafePerformIO' for randomness.
splitRand :: Double -> [a] -> ([a], [a])                                                            -- <TO-DO>: Does this need to be more sophisticated?
splitRand split xs = (take len shuffled, drop len shuffled)
  where
    len = ceiling (genericLength xs * split)
    shuffled = unsafePerformIO (shuffle xs)                                                         -- <TO-DO> 'unsafePerformIO'

-- | Naive-ish way to shuffle a list.
shuffle :: [a] -> IO [a]                                                                            -- <TO-DO>: Could shuffle in the top-level functions to avoid using unsafePerformIO?
shuffle x =                                                                                -- <TO-DO>: Does this need to be more sophisticated?
  if length x < 2
  then return x
  else do
   i <- randomRIO (0, length x - 1)
   r <- shuffle (take i x ++ drop (i + 1) x)
   return (x !! i : r)

-- | Compare the runtimes of two test programs, given as a list of
-- xy-coordinates, pointwise to generate improvement results. Only compare
-- results relating to test programs of the same airty. (Note: this should
-- always be the case when used in practice).
calculateImprovements
 :: ([(Double, Double)] -> Maybe (Ordering, Double))  -- '_improv' function from 'AnalOpts'.
 -> (Id, Either [Coord] [Coord3])                     -- Runtimes of test program 1.
 -> (Id, Either [Coord] [Coord3])                     -- Runtimes of test program 2.
 -> Maybe Improvement
calculateImprovements improv (idt1, Left cs1) (idt2, Left cs2) -- Unary against unary.
  -- If all coordinates didn't match, then can't generate improvements.
  | length toCompare < length cs1 = Nothing
  | otherwise = case improv toCompare of
      Nothing       -> Nothing
      Just (ord, d) -> Just (idt1, ord, idt2, d)
  where toCompare = matchCoords (sort cs1) (sort cs2)
calculateImprovements improv (idt1, Right cs1) (idt2, Right cs2) -- Binary against binary.
  -- If all coordinates didn't match, then can't generate improvements.
  | length toCompare < length cs1 = Nothing
  | otherwise = case improv toCompare of
      Nothing       -> Nothing
      Just (ord, d) -> Just (idt1, ord, idt2, d)
  where toCompare = matchCoords3 (sort cs1) (sort cs2)
calculateImprovements _ _ _ = Nothing -- Shouldn't happen.

-- | Make sure coordinates to compare when generating improvement results have
-- the same size information. (Note: this should always be the case because of
-- how they are generated, but better to check.). For 'Coord'.
matchCoords :: [Coord] -> [Coord] -> [(Double, Double)]
matchCoords [] _  = []
matchCoords _  [] = []
matchCoords ((s1, t1) : cs1) ((s1', t2) : cs2)
  | s1 == s1' = (t1, t2) : matchCoords cs1 cs2
  | otherwise = matchCoords cs1 cs2

-- | Make sure coordinates to compare when generating improvement results have
-- the same size information. (Note: this should always be the case because of
-- how they are generated, but better to check.). For 'Coord3'
matchCoords3 :: [Coord3] -> [Coord3] -> [(Double, Double)]
matchCoords3 [] _  = []
matchCoords3 _  [] = []
matchCoords3 ((s1, s2, t1) : cs1) ((s1', s2', t2) : cs2)
  | s1 == s1' && s2 == s2' = (t1, t2) : matchCoords3 cs1 cs2
  | otherwise = matchCoords3 cs1 cs2

-- | Fit the '_linearModels' in a given 'AnalOpts' to a data set to generate
-- 'LinearFit's (i.e., models with their missing parameters/coefficients
-- approximated). Currently this only works for 'Coord's, i.e., runtime
-- measurements from /unary/ test programs. Models are fitted to data
-- sets using ridge regression.
--
-- Note 'AnalOpts' is required here because it contains other settings needed
-- for fitting, i.e., cross-validation parameters: '_cvTrain', '_cvIters'.
--
-- Also sorts and filters the resulting models using '_statsSort' and
-- '_statsFilt' from 'AnalOpts', respectively.
fitCoords :: AnalOpts -> Either [Coord] [Coord3] -> [LinearFit]
fitCoords _ Right{} = []                                                                            -- <TO-DO>: Analyse 3D coords?
fitCoords aOpts (Left coords) = take (_topModels aOpts)
   $ sortBy (\lf1 lf2 -> (_statsSort aOpts) (_sts lf1) (_sts lf2))    -- SORT:   '_statsSort'
   $ filter (\lf -> (_statsFilt aOpts) (_sts lf))                     -- FILTER: '_statsFilt'
   $ catMaybes  -- The fitting process may fail in some cases, exactly when?                        -- <TO-DO>: When does fitting fail.. When not enough coords only!
   $ fmap ( ( candidateFit
               fitRidgeRegress           -- Use ridge regression to fit.
               (_cvTrain aOpts)          -- Train/evaluate data split.
               (_cvIters aOpts)          -- Number of cross-validation iterations.
               coords                    -- Data set.
            ) . generateLinearCandidate   -- 'LinearType' -> 'LinearCandidate'.
          ) (_linearModels aOpts)         -- Fit all models in 'AnalOpts'.


{-# LANGUAGE DataKinds #-}
module CannyEdgeEndPoint where

import           Control.Monad             as M
import           Data.Array.Repa           as R
import           Data.Binary               (decodeFile)
import           Data.ByteString           as BS (readFile)
import           Data.Complex
import           Data.List                 as L
import           DFT.Plan
import           FokkerPlanck.DomainChange
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           GHC.Word
import           Image.IO
import           Image.IO
import           Image.Transform
import           Utils.Parallel
import           OpenCV                    as CV hiding (Z)
import           Data.Array.Unboxed        as AU
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array

main = do
  args <- getArgs
  let (numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:histFilePath:numIterationStr:writeSourceFlagStr:cutoffRadiusEndPointStr:cutoffRadiusStr:reversalFactorStr:inputImgPath:threshold1Str:threshold2Str:pixelDistStr:patchNormFlagStr:patchNormSizeStr:approximatedEigenValueStr:numThreadStr:_) =
        args
      numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      scale0Freqs = [-scale0Freq .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      cutoffRadiusEndPoint = read cutoffRadiusEndPointStr :: Int
      cutoffRadius = read cutoffRadiusStr :: Int
      reversalFactor = read reversalFactorStr :: Double
      threshold1 = read threshold1Str :: Double
      threshold2 = read threshold2Str :: Double
      pixelDist = read pixelDistStr :: Int
      patchNormFlag = read patchNormFlagStr :: Bool
      patchNormSize = read patchNormSizeStr :: Int
      approximatedEigenValue = read approximatedEigenValueStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/CannyEdgeEndPoint"
  createDirectoryIfMissing True folderPath
  img <-
    (exceptError . coerceMat . imdecode ImreadUnchanged) <$>
    BS.readFile inputImgPath  :: IO (Mat ('S '[ 'D, 'D]) 'D ('S GHC.Word.Word8))
  let edge =
        exceptError . canny threshold1 threshold2 (Just 7) CannyNormL2 $ img
      edgeRepa = R.map fromIntegral . toRepa $ edge
      (Z :. _ :. rows :. cols) = extent edgeRepa
      nonzeros = L.filter (> 0) . R.toList $ edgeRepa
      avg = L.sum nonzeros / (fromIntegral . L.length $ nonzeros)
      edgeRepaSparse =
        computeS .
        pad [numPoint, numPoint, 1] 0 .
        R.extend (Z :. (1 :: Int) :. All :. All) .
        increasePixelDistance pixelDist .
        R.slice
          (R.map
             (\x ->
                if x >= avg
                  then 255
                  else 0)
             edgeRepa) $
        (Z :. (0 :: Int) :. All :. All)
  plotImageRepa (folderPath </> "edge.png") . ImageRepa 8 . computeS $ edgeRepa
  plotImageRepa (folderPath </> "edge_sparse.png") . ImageRepa 8 $
    edgeRepaSparse
  edgeRepaSparse <-
    (\(ImageRepa _ img) -> img) <$>
    readImageRepa (folderPath </> "edge_sparse.png") False
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          theta0Freqs
          thetaFreqs
          scale0Freqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqsp
             , L.length thetaFreqs
             ]
             0)
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (cutoff cutoffRadius radialArr)
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  -- planReversal <- makeR2Z2T0S0Plan emptyPlan arrR2Z2T0S0
  -- (plan, pathNormMethod) <-
  --   makePatchNormFilter
  --     planReversal
  --     numPoint
  --     numPoint
  --     patchNormFlag
  --     patchNormSize
  plan <- makeR2Z2T0S0Plan emptyPlan arrR2Z2T0S0
  pathNormMethod <-
    if patchNormFlag
      then do
        let points =
              createIndex2D .
              L.map fst . L.filter (\(_, v) -> v /= 0) . AU.assocs $
              (AU.listArray ((0, 0), (numPoint - 1, numPoint - 1)) . R.toList $
               edgeRepaSparse :: AU.Array (Int, Int) Double)
            ys =
              pointCluster
                (connectionMatrixP
                   (ParallelParams numThread 1)
                   (2 * pixelDist)
                   points) $
              points
        M.zipWithM_
          (\i ->
             plotImageRepa (folderPath </> (printf "Cluster%03d.png" i)) .
             ImageRepa 8)
          [1 :: Int ..] .
          cluster2Array numPoint numPoint $
          ys
        return . PowerMethodConnection $ ys
      else return PowerMethodGlobal
  let endPointFilePath =
        folderPath </>
        (printf
           "EndPoint_%d_%d_%d_%d_%d_%d_%.2f_%.2f_%f.dat"
           numPoint
           (round thetaFreq :: Int)
           (round scaleFreq :: Int)
           (round maxScale :: Int)
           (round tao :: Int)
           cutoffRadiusEndPoint
           thetaSigma
           scaleSigma
           reversalFactor)
      edgeRepaSparseNormalized = normalizeValueRange (0, 1) edgeRepaSparse
      bias =
        R.traverse
          edgeRepaSparseNormalized
          (const
             (Z :. (L.length theta0Freqs) :. (L.length scale0Freqs) :. numPoint :.
              numPoint)) $ \f (Z :. _ :. _ :. i :. j) ->
          f (Z :. (0 :: Int) :. i :. j) :+ 0
      eigenVec =
        R.traverse
          edgeRepaSparseNormalized
          (const
             (Z :. (L.length thetaFreqs) :. (L.length scaleFreqs) :.
              (L.length theta0Freqs) :.
              (L.length scale0Freqs) :.
              numPoint :.
              numPoint)) $ \f (Z :. k :. l :. _ :. _ :. i :. j) ->
          if k == div (L.length thetaFreqs) 2 &&
             l == div (L.length scaleFreqs) 2 &&
             f (Z :. (0 :: Int) :. i :. j) > 0
            then 1 /
                 (fromIntegral $ (L.length theta0Freqs * L.length scale0Freqs)) :+
                 0
            else 0
  -- plotImageRepaComplex (folderPath </> "bias_0.png") .
  --   ImageRepa 8 .
  --   computeS .
  --   R.extend (Z :. (1 :: Int) :. All :. All) .
  --   R.sumS . R.sumS . rotate4D . rotate4D $
  --   bias
  endPointFlag <- doesFileExist endPointFilePath
  endPointR2Z2 <-
    if endPointFlag
      then readRepaArray endPointFilePath
      else (do putStrLn "Couldn't find the endpoint data. Start computing..."
               arrR2Z2T0S0EndPoint <-
                 computeUnboxedP $
                 computeR2Z2T0S0ArrayRadial
                   (cutoff cutoffRadiusEndPoint radialArr)
                   numPoint
                   numPoint
                   1
                   maxScale
                   thetaFreqs
                   scaleFreqs
                   theta0Freqs
                   scale0Freqs
               completionFieldR2Z2 <-
                 computeS <$>
                 powerMethodR2Z2T0S0Reversal
                   plan
                   folderPath
                   numPoint
                   numPoint
                   numOrientation
                   thetaFreqs
                   theta0Freqs
                   numScale
                   scaleFreqs
                   scale0Freqs
                   maxScale
                   arrR2Z2T0S0EndPoint
                   pathNormMethod
                   numIteration
                   writeSourceFlag
                   (printf
                      "_%d_%d_%d_%d_%d_%d_%.2f_%.2f_%f_EndPoint"
                      numPoint
                      (round thetaFreq :: Int)
                      (round scaleFreq :: Int)
                      (round maxScale :: Int)
                      (round tao :: Int)
                      cutoffRadiusEndPoint
                      thetaSigma
                      scaleSigma
                      reversalFactor)
                   0.5
                   reversalFactor
                   approximatedEigenValue
                   (computeS bias)
                   eigenVec
               -- writeRepaArray endPointFilePath completionFieldR2Z2
               return completionFieldR2Z2)
  let endPointR2Z2Biased = R.zipWith (*) endPointR2Z2 bias
      biasMag =
        R.sumS .
        R.sumS .
        rotate4D .
        rotate4D .
        R.map magnitude .
        r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs $
        endPointR2Z2Biased
  plotImageRepa (folderPath </> "EndPointBias.png") .
    ImageRepa 8 . computeS . R.extend (Z :. (1 :: Int) :. All :. All) $
    biasMag
  -- printf "%f %f\n" reversalFactor (R.sumAllS biasMag)
  -- let rotatedEndBias =
  --       computeUnboxedS $
  --       R.zipWith
  --         (+)
  --         (rotateBiasR2Z2T0S0 90 thetaFreqs endPointR2Z2Biased)
  --         (rotateBiasR2Z2T0S0 (-90) thetaFreqs endPointR2Z2Biased)
  -- powerMethodR2Z2T0S0Bias
  --   plan
  --   folderPath
  --   numPoint
  --   numPoint
  --   numOrientation
  --   thetaFreqs
  --   theta0Freqs
  --   numScale
  --   scaleFreqs
  --   scale0Freqs
  --   arrR2Z2T0S0
  --   PowerMethodNone
  --   numIteration
  --   writeSourceFlag
  --   (printf
  --      "_%d_%d_%d_%d_%d_%d_%.2f_%.2f_%f"
  --      numPoint
  --      (round thetaFreq :: Int)
  --      (round scaleFreq :: Int)
  --      (round maxScale :: Int)
  --      (round tao :: Int)
  --      cutoffRadiusEndPoint
  --      thetaSigma
  --      scaleSigma
  --      reversalFactor)
  --   0.5
  --   rotatedEndBias
  --   eigenVec


module LatticeSymmetries.IO where

import Data.Aeson
import Data.Aeson.Types (typeMismatch)
import Data.Complex
import Data.Scientific

data SymmetrySpec = SymmetrySpec !(NonEmpty Int) !Int
  deriving stock (Read, Show, Eq)

instance FromJSON SymmetrySpec where
  parseJSON = withObject "symmetry" $ \v ->
    SymmetrySpec
      <$> v .: "permutation"
      <*> v .: "sector"

instance ToJSON SymmetrySpec where
  toJSON (SymmetrySpec permutation sector) =
    object ["permutation" .= permutation, "sector" .= sector]

data BasisSpec = BasisSpec !Int !(Maybe Int) !(Maybe Int) ![SymmetrySpec]
  deriving stock (Read, Show, Eq)

instance FromJSON BasisSpec where
  parseJSON = withObject "basis" $ \v ->
    BasisSpec
      <$> v .: "number_spins"
      <*> v .:? "hamming_weight"
      <*> v .:? "spin_inversion"
      <*> v .:! "symmetries" .!= []

instance ToJSON BasisSpec where
  toJSON (BasisSpec numberSpins hammingWeight spinInversion symmetries) =
    object
      [ "number_spins" .= numberSpins,
        "hamming_weight" .= maybe Null toJSON hammingWeight,
        "spin_inversion" .= maybe Null toJSON spinInversion,
        "symmetries" .= symmetries
      ]

newtype WrappedBasisSpec = WrappedBasisSpec BasisSpec

instance FromJSON WrappedBasisSpec where
  parseJSON = withObject "config" $ \v -> WrappedBasisSpec <$> v .: "basis"

instance FromJSON (Complex Double) where
  parseJSON (Number x) = pure . fromReal . toRealFloat $ x
    where
      fromReal :: Num a => a -> Complex a
      fromReal x' = x' :+ 0
  parseJSON v@(Array xs) = case (toList xs) of
    [re, im] -> (:+) <$> parseJSON re <*> parseJSON im
    _ -> typeMismatch "Complex" v
  parseJSON v = typeMismatch "Complex" v

newtype DenseMatrixSpec = DenseMatrixSpec [[Complex Double]]
  deriving stock (Read, Show, Eq)
  deriving newtype (FromJSON)

data InteractionSpec = InteractionSpec ![[Complex Double]] ![[Int]]
  deriving stock (Read, Show, Eq)

instance FromJSON InteractionSpec where
  parseJSON = withObject "interaction" $ \v ->
    InteractionSpec
      <$> v .: "matrix"
      <*> v .: "sites"

data OperatorSpec = OperatorSpec !Text !(NonEmpty InteractionSpec)
  deriving stock (Read, Show)

instance FromJSON OperatorSpec where
  parseJSON = withObject "operator" $ \v ->
    OperatorSpec
      <$> v .: "name"
      <*> v .: "terms"

data ConfigSpec = ConfigSpec !BasisSpec !OperatorSpec
  deriving stock (Read, Show)

instance FromJSON ConfigSpec where
  parseJSON = withObject "config" $ \v ->
    ConfigSpec
      <$> v .: "basis"
      <*> v .: "hamiltonian"


{-# LANGUAGE BangPatterns #-}
module Image.IO
  ( ImageRepa(..)
  , readImagePathList
  , imagePathSource
  , readImageRepa
  , readImageConduit
  , plotImageRepa
  , complexImageToColorImage
  , plotImageRepaComplex
  , plotImageRepaComplexGray
  ) where

import           Codec.Picture
import           Control.Monad.IO.Class       (liftIO)
import           Control.Monad.Trans.Resource
import           Data.Array.Repa              as R
import           Data.Complex
import           Data.Conduit                 as C
import           Data.Conduit.List            as CL
import           Data.List                    as L
import           Data.Vector.Unboxed          as VU
import           Data.Word
import           GHC.Float
import           Utils.Parallel

-- x = cols = width
-- y = rows = height
data ImageRepa a = ImageRepa
  { imageDepth   :: !Int
  , imageContent :: !(R.Array U DIM3 a) -- nf x cols x rows
  }

readImagePathList :: FilePath -> IO [String]
readImagePathList = fmap lines . readFile

imagePathSource :: FilePath -> ConduitT () FilePath (ResourceT IO) ()
imagePathSource filePath = do
  pathList <- liftIO $ readImagePathList filePath
  sourceList pathList

readImageRepa :: FilePath -> Bool -> IO (ImageRepa Double)
readImageRepa filePath isColor = do
  buffer <- liftIO $ readImage filePath
  case buffer of
    Left msg -> error msg
    Right dImg ->
      return $
      if isColor
        then case dImg of
               ImageY8 img ->
                 ImageRepa 8 . computeS $
                 fromFunction
                   (Z :. (1 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. _ :. x :. y) ->
                      fromIntegral $ pixelAt img x y :: Double)
               ImageY16 img ->
                 ImageRepa 16 . computeS $
                 fromFunction
                   (Z :. (1 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. _ :. x :. y) ->
                      fromIntegral $ pixelAt img x y :: Double)
               ImageYF img ->
                 ImageRepa 64 . computeS $
                 fromFunction
                   (Z :. (1 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. _ :. x :. y) ->
                      float2Double $ pixelAt img x y :: Double)
               ImageRGB8 img ->
                 ImageRepa 8 . computeS $
                 fromFunction
                   (Z :. (3 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. k :. x :. y) ->
                      let (PixelRGB8 r g b) = pixelAt img x y
                       in case k of
                            0 -> fromIntegral r
                            1 -> fromIntegral g
                            2 -> fromIntegral b
                            _ -> error "readImageConduit: dimension error.")
               ImageRGB16 img ->
                 ImageRepa 16 . computeS $
                 fromFunction
                   (Z :. (3 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. k :. x :. y) ->
                      let (PixelRGB16 r g b) = pixelAt img x y
                       in case k of
                            0 -> fromIntegral r
                            1 -> fromIntegral g
                            2 -> fromIntegral b
                            _ -> error "readImageConduit: dimension error.")
               ImageRGBF img ->
                 ImageRepa 64 . computeS $
                 fromFunction
                   (Z :. (3 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. k :. x :. y) ->
                      let (PixelRGBF r g b) = pixelAt img x y
                       in case k of
                            0 -> float2Double r
                            1 -> float2Double g
                            2 -> float2Double b
                            _ -> error "readImageConduit: dimension error.")
               img ->
                 let rgbImg = convertRGB8 img
                  in ImageRepa 8 . computeS $
                     fromFunction
                       (Z :. (3 :: Int) :. imageWidth rgbImg :.
                        imageHeight rgbImg)
                       (\(Z :. k :. x :. y) ->
                          let (PixelRGB8 r g b) = pixelAt rgbImg x y
                           in case k of
                                0 -> fromIntegral r
                                1 -> fromIntegral g
                                2 -> fromIntegral b
                                _ -> error "readImageConduit: dimension error.")
        else case dImg of
               ImageY8 img ->
                 ImageRepa 8 . computeS $
                 fromFunction
                   (Z :. (1 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. _ :. x :. y) ->
                      fromIntegral $ pixelAt img x y :: Double)
               ImageY16 img ->
                 ImageRepa 16 . computeS $
                 fromFunction
                   (Z :. (1 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. _ :. x :. y) ->
                      fromIntegral $ pixelAt img x y :: Double)
               ImageYF img ->
                 ImageRepa 64 . computeS $
                 fromFunction
                   (Z :. (1 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. _ :. x :. y) ->
                      float2Double $ pixelAt img x y :: Double)
               ImageRGB8 img ->
                 ImageRepa 8 . computeS $
                 fromFunction
                   (Z :. (1 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. _ :. x :. y) ->
                      let (PixelRGB8 r g b) = pixelAt img x y
                       in rgb2Gray
                            (fromIntegral (maxBound :: Word8))
                            (fromIntegral r)
                            (fromIntegral g)
                            (fromIntegral b))
               ImageRGB16 img ->
                 ImageRepa 16 . computeS $
                 fromFunction
                   (Z :. (1 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. _ :. x :. y) ->
                      let (PixelRGB16 r g b) = pixelAt img x y
                       in rgb2Gray
                            (fromIntegral (maxBound :: Word16))
                            (fromIntegral r)
                            (fromIntegral g)
                            (fromIntegral b))
               ImageRGBF img ->
                 ImageRepa 64 . computeS $
                 fromFunction
                   (Z :. (1 :: Int) :. imageWidth img :. imageHeight img)
                   (\(Z :. _ :. x :. y) ->
                      let (PixelRGBF r g b) = pixelAt img x y
                       in rgb2Gray
                            1
                            (float2Double r)
                            (float2Double g)
                            (float2Double b))
               img ->
                 let rgbImg = convertRGB8 img
                  in ImageRepa 8 . computeS $
                     fromFunction
                       (Z :. (1 :: Int) :. imageWidth rgbImg :.
                        imageHeight rgbImg)
                       (\(Z :. _ :. x :. y) ->
                          let (PixelRGB8 r g b) = pixelAt rgbImg x y
                           in rgb2Gray
                                (fromIntegral (maxBound :: Word8))
                                (fromIntegral r)
                                (fromIntegral g)
                                (fromIntegral b))

readImageConduit ::
     Bool -> ConduitT FilePath (ImageRepa Double) (ResourceT IO) ()
readImageConduit isColor =
  awaitForever
    (\filePath -> do
       img <- liftIO $ readImageRepa filePath isColor
       yield img)

{-# INLINE rgb2Gray #-}
rgb2Gray :: Double -> Double -> Double -> Double -> Double
rgb2Gray bound r g b
  | yLinear <= 0.0031308 = 12.92 * yLinear * bound
  | otherwise = (1.055 * (yLinear ** (1 / 2.4)) - 0.055) * bound
  where
    yLinear =
      0.2126 * func bound r + 0.7152 * func bound g + 0.0722 * func bound b

{-# INLINE func #-}
func :: Double -> Double -> Double
func bound x
  | y < 0.04045 = y / 12.92
  | otherwise = ((y + 0.055) / 1.055) ** 2.4
  where
    y = x / bound

{-# INLINE normalizeImageRepa #-}
normalizeImageRepa :: (RealFrac e, Eq e, Unbox e, Ord e) => ImageRepa e -> ImageRepa e
normalizeImageRepa i@(ImageRepa depth img)
  | maxV == minV =
    ImageRepa depth . computeS . R.map (\x -> fromIntegral (2 ^ depth - 1)) $
    img
  | otherwise =
    ImageRepa depth .
    computeS .
    R.map (\x -> (x - minV) / (maxV - minV) * fromIntegral (2 ^ depth - 1)) $
    img
  where
    !maxV = VU.maximum . toUnboxed $ img
    !minV = VU.minimum . toUnboxed $ img

{-# INLINE normalizeImageRepaPar #-}
normalizeImageRepaPar :: (RealFrac e, Eq e, Unbox e, Ord e) => ImageRepa e -> IO (ImageRepa e)
normalizeImageRepaPar i@(ImageRepa depth img) = do
  if maxV == minV
    then do
      arr <- computeP . R.map (\x -> fromIntegral (2 ^ depth - 1)) $img
      return . ImageRepa depth $ arr
    else do
      arr <-
        computeP .
        R.map (\x -> (x - minV) / (maxV - minV) * fromIntegral (2 ^ depth - 1)) $
        img
      return . ImageRepa depth $ arr
  where
    !maxV = VU.maximum . toUnboxed $ img
    !minV = VU.minimum . toUnboxed $ img

{-# INLINE plotImageRepa #-}
plotImageRepa :: (Unbox e, Ord e, RealFrac e) => FilePath -> ImageRepa e -> IO ()
plotImageRepa filePath img@(ImageRepa depth x) = do
  normalizedImg <- imageContent <$> normalizeImageRepaPar img
  let Z :. nfp' :. nxp' :. nyp' = extent x
      -- normalizedImg = imageContent . normalizeImageRepa $ img
      w =
        case nfp' of
          1
            -- ImageY8 $
            -- generateImage
            --   (\x y ->
            --      let v =
            --            fromIntegral . round $
            --            normalizedImg R.! (Z :. 0 :. x :. y)
            --       in v)
            --   nxp'
            --   nyp'
           ->
            ImageRGB8 $
            generateImage
              (\x y ->
                 let r =
                       fromIntegral . round $
                       normalizedImg R.! (Z :. 0 :. x :. y)
                     g =
                       fromIntegral . round $
                       normalizedImg R.! (Z :. 0 :. x :. y)
                     b =
                       fromIntegral . round $
                       normalizedImg R.! (Z :. 0 :. x :. y)
                  in PixelRGB8 r g b)
              nxp'
              nyp'
          3 ->
            ImageRGB8 $
            generateImage
              (\x y ->
                 let r =
                       fromIntegral . round $
                       normalizedImg R.! (Z :. 0 :. x :. y)
                     g =
                       fromIntegral . round $
                       normalizedImg R.! (Z :. 1 :. x :. y)
                     b =
                       fromIntegral . round $
                       normalizedImg R.! (Z :. 2 :. x :. y)
                  in PixelRGB8 r g b)
              nxp'
              nyp'
          _ ->
            error $
            "ImageRepa is neither a gray image nor a color image. There are " L.++
            show nfp' L.++
            " channels."
  savePngImage filePath w

-- Plot Complex Image
{-# INLINE complexScale #-}
complexScale :: (Ord e, Floating e, Unbox e) => ImageRepa (Complex e) -> e
complexScale (ImageRepa _ arr) =
  let (rVec, iVec) = VU.unzip . VU.map (\(a :+ b) -> (a, b)) . R.toUnboxed $ arr
      rMax = VU.maximum rVec
      rMin = VU.minimum rVec
      iMax = VU.maximum iVec
      iMin = VU.minimum iVec
   in 2 / ((max rMax iMax) - (min rMin iMin))

{-# INLINE complexImageToColorImage #-}
complexImageToColorImage ::
     (Floating e, Ord e, Unbox e) => ImageRepa (Complex e) -> ImageRepa e
complexImageToColorImage img@(ImageRepa depth arr) =
  let scale = complexScale img
      (Z :. _ :. cols :. rows) = extent arr
      (r, g, b) =
        VU.unzip3 .
        VU.map
          (\(r :+ i) ->
             let [x, y] = L.map (* scale) [r, i]
                 radius = sqrt (x * x + y * y)
                 a = 0.40824829046386301636 * x
                 b = 0.70710678118654752440 * y
                 d = 1 / (1 + radius * radius)
                 d' = 0.5 - radius * d
                 r' = 0.5 + (0.81649658092772603273 * x * d)
                 b' = 0.5 - (d * (a - b))
                 g' = 0.5 - (d * (a + b))
                 [red, grn, blu] = L.map (+ d') [r', g', b']
                 [red', grn', blu'] = L.map (flip (-) d') [r', g', b']
              in if radius < 1
                   then (red', grn', blu')
                   else (red, grn, blu)) .
        toUnboxed $
        arr
   in ImageRepa 8 . fromUnboxed (Z :. (3 :: Int) :. cols :. rows) . VU.concat $
      [r, g, b]

{-# INLINE normalizeComplexImage #-}
normalizeComplexImage ::
     ImageRepa (Complex Double) -> ImageRepa (Complex Double)
normalizeComplexImage (ImageRepa depth img) =
  let maxV = VU.maximum magVec
      minV = VU.minimum magVec
      maxVal = 2 ^ depth
      minVal = 0
      magVec = VU.map magnitude . toUnboxed $ img
   in if maxV == minV
        then (ImageRepa depth img)
        else ImageRepa depth .
             fromUnboxed (extent img) .
             VU.zipWith (\a b -> mkPolar b (phase a)) (toUnboxed img) .
             VU.map
               (\x -> (x - minV) / (maxV - minV) * (maxVal - minVal) + minVal) $
             magVec

{-# INLINE plotImageRepaComplex #-}
plotImageRepaComplex ::
     (Unbox e, Floating e, RealFrac e)
  => FilePath
  -> ImageRepa (Complex e)
  -> IO ()
plotImageRepaComplex filePath img =
  plotImageRepa filePath . complexImageToColorImage -- . normalizeComplexImage
  $
  img

{-# INLINE plotImageRepaComplexGray #-}
plotImageRepaComplexGray :: FilePath -> ImageRepa (Complex Double) -> IO ()
plotImageRepaComplexGray filePath img =
  let (ImageRepa depth arr) = complexImageToColorImage img
      (Z :. chs :. cols :. rows) = extent arr
  in plotImageRepa filePath .
     ImageRepa depth .
     computeS .
     extend (Z :. (1 :: Int) :. All :. All) .
     R.sumS .
     R.backpermute
       (Z :. cols :. rows :. chs)
       (\(Z :. c :. r :. k) -> (Z :. k :. c :. r)) $
     arr


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}

module HVX.Primitives
-- TODO(mh): Currently primitives that should generate implicit constraints are
-- unsupported. The should be supported shortily. (2014-06-04)
  ( apply
  , hadd
  , (+~)
  , hmul
  , (*~)
  , habs
  , neg
  , hlog
  , hexp
  , logsumexp
  , hmax
  , hmin
  , norm
  , berhu
  , huber
  , quadform
  , powBaseP0
  , powBaseP01
  , powBaseP1
  , powBaseP1InfEven
  , powBaseP1InfNotInt
  ) where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.HMatrix (eigSH')

import HVX.Internal.DCP
import HVX.Internal.Matrix
import HVX.Internal.Primitives
import HVX.Internal.Util

apply :: (ApplyVex vf mf ve me ~ vr, ValidVex vr)
  => Fun vf mf -> Expr ve me -> Expr vr (ApplyMon mf me)
apply = EFun

hadd :: (AddVex v1 v2 ~ v3, ValidVex v3)
  => Expr v1 m1 -> Expr v2 m2 -> Expr v3 (AddMon m1 m2)
hadd = EAdd

infixl 6 +~
(+~) :: (AddVex v1 v2 ~ v3, ValidVex v3)
  => Expr v1 m1 -> Expr v2 m2 -> Expr v3 (AddMon m1 m2)
(+~) = hadd

-- Constructors that enforce DCP constraints.
hmul :: (ApplyVex 'Affine 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Expr 'Affine 'Const -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
hmul (EConst a) e = apply (Mul a) e
hmul _ _ = error "the left argument of a multiply must be a constant"

infixl 7 *~
(*~) :: (ApplyVex 'Affine 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Expr 'Affine 'Const -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
(*~) (EConst a) e = apply (Mul a) e
(*~) _ _ = error "the left argument of a multiply must be a constant"

habs :: (ApplyVex 'Convex 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
habs = apply Abs

neg :: (ApplyVex 'Affine 'Noninc v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Noninc m1)
neg = apply Neg

hlog :: (ApplyVex 'Concave 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
hlog = apply Log

hexp :: (ApplyVex 'Convex 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
hexp e = apply Exp e

logsumexp :: (ApplyVex 'Convex 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
logsumexp = apply LogSumExp

hmax :: (ApplyVex 'Convex 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
hmax = apply Max

hmin :: (ApplyVex 'Concave 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
hmin = apply Min

norm :: (ApplyVex 'Convex 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
norm p
  | p == infinity = error "Internal: Infinity norm should become max . abs."
  | 1 <= p = apply (Norm p)
  | otherwise = error "Internal: Norm only supports p >= 1."

berhu :: (ApplyVex 'Convex 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
berhu m
  | 0 < m = apply (Berhu m)
  | otherwise = error "Internal: Berhu only supports m >= 0."

huber :: (ApplyVex 'Convex 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
huber m
  | 0 < m = apply (Huber m)
  | otherwise = error "Internal: Huber only supports m >= 0."

quadform :: (ApplyVex 'Convex 'Nonmon 'Affine m1 ~ v2, ValidVex v2)
  => Expr 'Affine 'Const -> Expr 'Affine m1 -> Expr v2 (ApplyMon 'Nonmon m1)
quadform (EConst a) e
  | rows a == cols a
    && fpequalsMat a (tr a)
    && 0 <= (maxElement.fst.eigSH' $ a) = apply (Quadform a) e -- we have checked that the matrix being passed in is Hermitian, so it's safe to use eigSH'.
  | otherwise = error "Matrices in quadratic forms must be positive semidefinite."
quadform _ _ = error "The matrix sandwitched by the quadratic form must be constant."

powBaseP0 :: (ApplyVex 'Affine 'Const v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Const m1)
powBaseP0 p
  | p `fpequals` 0 = apply PowBaseP0
  | otherwise = error "Internal: PowBaseP0 only supports p == 0."

powBaseP01 :: (ApplyVex 'Concave 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
powBaseP01 p
  | 0 < p && p < 1 = apply (PowBaseP01 p)
  | otherwise = error "Internal: PowBaseP01 only supports 0 < p < 1."

powBaseP1 :: (ApplyVex 'Affine 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
powBaseP1 p
  | p `fpequals` 1 = apply PowBaseP1
  | otherwise = error "Internal: PowBaseP1 only supports p == 1."

powBaseP1InfEven :: (ApplyVex 'Convex 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
powBaseP1InfEven p
  | 1 < p && isInteger p && even intP = apply (PowBaseP1InfEven intP)
  | otherwise = error "Internal: PowBaseP1InfEven only supports even p > 1."
  where intP = round p :: Integer

powBaseP1InfNotInt :: (ApplyVex 'Convex 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
powBaseP1InfNotInt p
  | 1 < p && not (isInteger p) = apply (PowBaseP1InfNotInt p)
  | otherwise = error "Internal: PowBaseP1InfNotInt only supports non integral p > 1."


-- |
-- Module      :  Acme.Operators.Base
-- License     :  BSD-3-Clause (see the file LICENSE)
-- Maintainer  :  Oleg Grenrus <oleg.grenrus@iki.fi>
--
-- Have you ever been wondering about some magic lookin operator. Here are all operators in base package. Not so many.
--
-- The operators are in few bigger logical groups, subgrouped by package. Many of them are exported from 'Prelude'.
--
-- There is a Stack Overflow QA listing prononciations of some of the operators.
-- See: <http://stackoverflow.com/questions/7746894/are-there-pronounceable-names-for-common-haskell-operators>
{-# LANGUAGE CPP #-}
module Acme.Operators.Base (
  -- * Functions
  -- ** Data.Function
  ($),
  (&),
  -- * Functors
  -- | Infamous operators on 'Functor', 'Applicative' and 'Monad'. Also 'Alternative'.

  -- ** Data.Functor
  (<$>),
  (<$),
  ($>),
  -- ** Control.Applicative
  (<*>),
  (<*),
  (*>),
  (<**>),
  (<|>),
  -- ** Control.Monad
  (>>=),
  (>>),
  (=<<),
  (>=>),
  (<=<),
  (<$!>),
  -- * Category & Arrows
  -- ** Control.Category
  -- | Dot operator '.' clashes with one in 'Prelude': 'Prelude..'. You can use '<<<'.
  (.),
  (>>>),
  (<<<),
  -- ** Control.Arrow
  -- | Optionally one can use @Arrows@ language extensions. See: <https://downloads.haskell.org/~ghc/7.10.1/docs/html/users_guide/arrow-notation.html>.
  (***),
  (&&&),
  (^>>),
  (>>^),
  (<<^),
  (^<<),
  (+++),
  (|||),
  -- * Algebraic
  -- ** Data.Eq
  (==),
  (/=),
  -- ** Data.Ord
  (<=),
  (<),
  (>),
  (>=),
  -- * Data.Monoid
  (<>),
  -- * Data
  -- ** Data.Bits
  (.&.),
  (.|.),
  -- ** Data.Bool
  (&&),
  (||),
  -- ** Data.Complex
  Complex(..), -- :+
  -- ** Data.Ratio
  (%),
  -- ** Data.List
  (++),
  (!!),
  (\\),
  -- * Numeric
  (+),
  (*),
  (-),
  (/),
  -- ** Expontentials
  (^),
  (^^),
  (**),
  -- * Type level trickery
  -- ** Data.Type.Equality
  (:~:),
  ) where

import Control.Applicative
import Control.Arrow
import Control.Category
import Control.Monad
import Data.Bits
import Data.Complex (Complex(..))
import Data.Function hiding ((.))
import Data.Functor
import Data.List
import Data.Monoid
import Data.Ratio
import Data.Type.Equality
import Prelude hiding ((.))

#if !MIN_VERSION_base(4,8,0)

infixl 4 <$!>

-- | Strict version of 'Data.Functor.<$>'.
--
-- @since 4.8.0.0
(<$!>) :: Monad m => (a -> b) -> m a -> m b
{-# INLINE (<$!>) #-}
f <$!> m = do
  x <- m
  let z = f x
  z `seq` return z

infixl 1 &

-- | '&' is a reverse application operator.  This provides notational
-- convenience.  Its precedence is one higher than that of the forward
-- application operator '$', which allows '&' to be nested in '$'.
--
-- @since 4.8.0.0
(&) :: a -> (a -> b) -> b
x & f = f x

#endif


{-# LANGUAGE CPP                  #-}
{-# LANGUAGE DeriveDataTypeable   #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE NoImplicitPrelude    #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE StandaloneDeriving   #-}
{-# LANGUAGE UndecidableInstances #-}

--------------------------------------------------------------------------------
-- |
-- Module:     Numeric.LevMar
-- Copyright:  (c) 2009 - 2014 Roel van Dijk & Bas van Dijk
-- License:    BSD-style (see the file LICENSE)
-- Maintainer: Roel van Dijk <vandijk.roel@gmail.com>
--             Bas van Dijk <v.dijk.bas@gmail.com>
-- Stability:  Experimental
--
-- For additional documentation see the
-- <http://www.ics.forth.gr/~lourakis/levmar/ documentation of the levmar C>
-- library which this library is based on:
--
--------------------------------------------------------------------------------

module Numeric.LevMar
    ( -- * Model & Jacobian.
      Params, Samples
    , Model, Jacobian

      -- * Levenberg-Marquardt algorithm.
    , LevMarable(levmar)

      -- * Minimization options.
    , Options(..), defaultOpts

      -- * Constraints
    , Constraints(..), LinearConstraints

      -- * Output
    , Info(..), StopReason(..), LevMarError(..)
    ) where

-------------------------------------------------------------------------------
-- Imports
-------------------------------------------------------------------------------

-- from base:
import Control.Monad         ( return, mplus )
import Control.Exception     ( Exception )
import Data.Bool             ( (&&), (||), otherwise )
import Data.Data             ( Data )
import Data.Typeable         ( Typeable )
import Data.Either           ( Either(Left, Right) )
import Data.Eq               ( Eq, (==), (/=) )
import Data.Function         ( (.), ($) )
import Data.Functor          ( (<$>) )
import Data.Int              ( Int )
import Data.List             ( lookup, (++) )
import Data.Maybe            ( Maybe(Nothing, Just), isJust, fromJust, fromMaybe )
import Data.Monoid           ( Monoid, mempty, mappend )
import Data.Ord              ( Ord, (<) )
import Foreign.C.Types       ( CInt )
import Foreign.Marshal.Array ( allocaArray, withArray, peekArray, copyArray )
import Foreign.Ptr           ( Ptr, nullPtr )
import Foreign.ForeignPtr    ( ForeignPtr, newForeignPtr_, withForeignPtr )
import Foreign.Storable      ( Storable )
import Prelude               ( Num, Enum, Fractional, RealFrac, Float, Double
                             , fromIntegral, toEnum, (-), (*), error, floor
                             )
import System.IO             ( IO )
import System.IO.Unsafe      ( unsafePerformIO )
import Text.Read             ( Read )
import Text.Show             ( Show, show )

#if __GLASGOW_HASKELL__ >= 605
import GHC.ForeignPtr        ( mallocPlainForeignPtrBytes )
import Prelude               ( undefined )
import Foreign.Storable      ( sizeOf )
#else
import Foreign.ForeignPtr    ( mallocForeignPtrArray )
#endif

#if __GLASGOW_HASKELL__ < 700
import Prelude               ( fromInteger, (>>=), (>>), fail )
#endif

-- from hmatrix:
#if MIN_VERSION_hmatrix(0,17,0)
import Numeric.LinearAlgebra.Data ( Matrix, flatten, rows, reshape )
import Numeric.LinearAlgebra      ( Container, Element )
#else
import Data.Packed.Matrix         ( Matrix, Element, flatten, rows, reshape )
import Numeric.Container          ( Container )
import Numeric.LinearAlgebra      ( {- Instances for Matrix -} )
#endif

-- from vector:
import           Data.Vector.Storable       ( Vector )
import qualified Data.Vector.Storable as VS ( unsafeWith, length
                                            , unsafeFromForeignPtr
                                            , length
                                            )

-- from bindings-levmar:
import Bindings.LevMar ( c'LM_INFO_SZ

                       , withModel
                       , withJacobian

                       , c'LM_ERROR
                       , c'LM_ERROR_LAPACK_ERROR
                       , c'LM_ERROR_FAILED_BOX_CHECK
                       , c'LM_ERROR_MEMORY_ALLOCATION_FAILURE
                       , c'LM_ERROR_CONSTRAINT_MATRIX_ROWS_GT_COLS
                       , c'LM_ERROR_CONSTRAINT_MATRIX_NOT_FULL_ROW_RANK
                       , c'LM_ERROR_TOO_FEW_MEASUREMENTS
                       , c'LM_ERROR_SINGULAR_MATRIX
                       , c'LM_ERROR_SUM_OF_SQUARES_NOT_FINITE

                       , c'LM_INIT_MU
                       , c'LM_STOP_THRESH
                       , c'LM_DIFF_DELTA
                       )
import qualified Bindings.LevMar ( Model, Jacobian )

-- from levmar (this package):
import Bindings.LevMar.CurryFriendly ( LevMarDer,     LevMarDif
                                     , LevMarBCDer,   LevMarBCDif
                                     , LevMarLecDer,  LevMarLecDif
                                     , LevMarBLecDer, LevMarBLecDif

                                     , dlevmar_der,      slevmar_der
                                     , dlevmar_dif,      slevmar_dif
                                     , dlevmar_bc_der,   slevmar_bc_der
                                     , dlevmar_bc_dif,   slevmar_bc_dif
                                     , dlevmar_lec_der,  slevmar_lec_der
                                     , dlevmar_lec_dif,  slevmar_lec_dif
                                     , dlevmar_blec_der, slevmar_blec_der
                                     , dlevmar_blec_dif, slevmar_blec_dif
                                     )

--------------------------------------------------------------------------------
-- Model & Jacobian.
--------------------------------------------------------------------------------

-- | Parameter vector of length @m@.
--
-- Ensure that @m <= n@ where @n@ is the length of the 'Samples' vector.
type Params r = Vector r

-- | Sample vector of length @n@.
--
-- Ensure that @n >= m@ where @m@ is the length of the 'Params' vector.
type Samples r = Vector r

{-| A functional relation describing measurements represented as a function
from a vector of parameters to a vector of expected samples.

 * Ensure that the length @m@ of the parameter vector equals the length of the
   initial parameter vector in 'levmar'.

 * Ensure that the length @n@ of the output sample vector equals the length of
   the sample vector in 'levmar'.

 * Ensure that the length @n@ of the output sample vector vector is bigger than or
   equal to the length @m@ of the parameter vector.
-}
type Model r = Params r -> Samples r

{-| The <http://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant jacobian>
of the 'Model' function. Expressed as a function from a vector of
parameters to a matrix which for each expected sample describes the
partial derivatives of the parameters.

 * Ensure that the length @m@ of the parameter vector equals the length of the
   initial parameter vector in 'levmar'.

 * Ensure that the output matrix has the dimension @n><m@ where @n@ is the
   number of samples and @m@ is the number of parameters.
-}
type Jacobian r = Params r -> Matrix r

--------------------------------------------------------------------------------
-- Levenberg-Marquardt algorithm.
--------------------------------------------------------------------------------

-- | The Levenberg-Marquardt algorithm is overloaded to work on 'Double' and 'Float'.
class LevMarable r where

    -- | The Levenberg-Marquardt algorithm.
    --
    -- Returns a triple of the found parameters, a structure containing
    -- information about the minimization and the covariance matrix
    -- corresponding to LS solution.
    --
    -- Ensure that @n >= m@.
    levmar :: Model r            -- ^ Model
           -> Maybe (Jacobian r) -- ^ Optional jacobian
           -> Params r           -- ^ Initial parameters of length @m@
           -> Samples r          -- ^ Sample vector of length @n@
           -> Int                -- ^ Maximum iterations
           -> Options r          -- ^ Minimization options
           -> Constraints r      -- ^ Constraints
           -> Either LevMarError (Params r, Info r, Matrix r)

instance LevMarable Float where
    levmar = gen_levmar slevmar_der
                        slevmar_dif
                        slevmar_bc_der
                        slevmar_bc_dif
                        slevmar_lec_der
                        slevmar_lec_dif
                        slevmar_blec_der
                        slevmar_blec_dif

instance LevMarable Double where
    levmar = gen_levmar dlevmar_der
                        dlevmar_dif
                        dlevmar_bc_der
                        dlevmar_bc_dif
                        dlevmar_lec_der
                        dlevmar_lec_dif
                        dlevmar_blec_der
                        dlevmar_blec_dif

{-| @gen_levmar@ takes the low-level C functions as arguments and
executes one of them depending on the optional jacobian and constraints.

Preconditions:

@
  length ys >= length ps

     isJust mLowBs && length (fromJust mLowBs) == length ps
  && isJust mUpBs  && length (fromJust mUpBs)  == length ps

  boxConstrained && (all $ zipWith (<=) (fromJust mLowBs) (fromJust mUpBs))
@
-}
gen_levmar :: forall r. (RealFrac r, Element r)
           => LevMarDer r
           -> LevMarDif r
           -> LevMarBCDer r
           -> LevMarBCDif r
           -> LevMarLecDer r
           -> LevMarLecDif r
           -> LevMarBLecDer r
           -> LevMarBLecDif r

           -> Model r            -- ^ Model
           -> Maybe (Jacobian r) -- ^ Optional jacobian
           -> Params r           -- ^ Initial parameters
           -> Samples r          -- ^ Samples
           -> Int                -- ^ Maximum iterations
           -> Options r          -- ^ Options
           -> Constraints r      -- ^ Constraints
           -> Either LevMarError (Params r, Info r, Matrix r)
gen_levmar f_der
           f_dif
           f_bc_der
           f_bc_dif
           f_lec_der
           f_lec_dif
           f_blec_der
           f_blec_dif
           model mJac ps ys itMax opts (Constraints mLowBs mUpBs mWeights mLinC)
               | m == 0    = Left LevMarError -- LAPACK will crash otherwise!
               | otherwise =
  -- All effects are contained, so we can safely perform:
  unsafePerformIO $ do

    -- We need to pass the initial parameters 'ps' to the C function.
    -- However, we can't just pass a pointer to them because the C function
    -- will modify the parameters during execution which will violate
    -- referential transparanency. Instead we allocate new space
    -- and copy the parameters to it.
    --
    -- Note that, in the end, the array is returned from this function.
    -- This means that the only way to guarantee its finalisation
    -- is to allocate it using a ForeignPtr:
    psFP <- fastMallocForeignPtrArray m
    withForeignPtr psFP $ \psPtr -> do
      VS.unsafeWith ps $ \psPtrInp ->
        copyArray psPtr psPtrInp m

      -- Retrieve the (read-only) pointer 'ysPtr' to the samples vector 'ys'
      -- so we can pass it to the C function:
      VS.unsafeWith ys $ \ysPtr ->

        -- Convert the Options 'opts' to a list and then to an array
        -- so we can pass the (read-only) pointer 'optsPtr' to the C function:
        withArray (optsToList opts) $ \optsPtr ->

          -- Allocate space for the info array
          -- so we can pass it to the C function.
          -- Note that, in the end, this array is converted to an Info value
          -- and returned from this function.
          allocaArray c'LM_INFO_SZ $ \infoPtr -> do

            -- Allocate space for the covariance matrix
            -- so we can pass it to the C function.
            -- Like the parameters array the matrix
            -- needs to be returned from this function.
            -- So we also allocate it using a ForeignPtr:
            covarFP <- fastMallocForeignPtrArray mm
            withForeignPtr covarFP $ \covarPtr ->

              -- 'cmodel' is the low-level model function which is converted
              -- to the FunPtr 'modelFunPtr' and passed to the C function.
              -- 'cmodel' will first convert the parameters pointer 'parPtr'
              -- into a Vector after converting it into a ForeignPtr
              -- (without a finalizer).
              -- Then it will apply the high-level 'model' function
              -- to this parameter vector. The resulting vector is then copied
              -- to the output buffer 'hxPtr':
              let cmodel :: Bindings.LevMar.Model r
                  cmodel parPtr hxPtr _ _ _ = do
                    parFP <- newForeignPtr_ parPtr
                    let psV = VS.unsafeFromForeignPtr parFP 0 m
                        vector = model psV
                    VS.unsafeWith vector $ \p -> copyArray hxPtr p (VS.length vector)
              in withModel cmodel $ \modelFunPtr -> do

                 -- All the low-level C functions share a common set of arguments.
                 -- 'runDif' applies these arguments to the given C function 'f':
                 let runDif :: LevMarDif r -> IO CInt
                     runDif f = f modelFunPtr
                                  psPtr
                                  ysPtr
                                  (fromIntegral m)
                                  (fromIntegral n)
                                  (fromIntegral itMax)
                                  optsPtr
                                  infoPtr
                                  nullPtr
                                  covarPtr
                                  nullPtr

                 err <- case mJac of
                   Nothing -> if boxConstrained
                              then if linConstrained
                                   then withBoxConstraints
                                          (withLinConstraints $ withWeights runDif)
                                          f_blec_dif
                                   else withBoxConstraints runDif f_bc_dif
                              else if linConstrained
                                   then withLinConstraints runDif f_lec_dif
                                   else runDif f_dif

                   Just jac ->
                     let cjacobian :: Bindings.LevMar.Jacobian r
                         cjacobian parPtr jPtr _ _ _ = do
                           parFP <- newForeignPtr_ parPtr
                           let psV    = VS.unsafeFromForeignPtr parFP 0 m
                               matrix = jac psV
                               vector = flatten matrix
                           VS.unsafeWith vector $ \p ->
                             copyArray jPtr p (VS.length vector)
                     in withJacobian cjacobian $ \jacobPtr ->

                       let runDer :: LevMarDer r -> IO CInt
                           runDer f = runDif $ f jacobPtr
                       in if boxConstrained
                          then if linConstrained
                               then withBoxConstraints
                                      (withLinConstraints $ withWeights runDer)
                                      f_blec_der
                               else withBoxConstraints runDer f_bc_der
                          else if linConstrained
                               then withLinConstraints runDer f_lec_der
                               else runDer f_der

                 -- Handling errors:
                 if err < 0
                    -- we don't treat the following two as an error:
                    && err /= c'LM_ERROR_SINGULAR_MATRIX
                    && err /= c'LM_ERROR_SUM_OF_SQUARES_NOT_FINITE
                   then return $ Left $ convertLevMarError err

                   else do -- Converting results:
                     info <- listToInfo <$> peekArray c'LM_INFO_SZ infoPtr
                     let psV = VS.unsafeFromForeignPtr psFP 0 m
                     let covarM = reshape m $ VS.unsafeFromForeignPtr covarFP 0 mm

                     return $ Right (psV, info, covarM)
  where
    m  = VS.length ps
    n  = VS.length ys
    mm = m*m

    -- Whether the parameters are constrained by a linear equation.
    linConstrained = isJust   mLinC
    (cMat, rhcVec) = fromJust mLinC

    -- Whether the parameters are constrained by a bounding box.
    boxConstrained = isJust mLowBs || isJust mUpBs

    withBoxConstraints f g =
        maybeWithArray mLowBs $ \lBsPtr ->
          maybeWithArray mUpBs $ \uBsPtr ->
            f $ g lBsPtr uBsPtr

    withLinConstraints f g =
        VS.unsafeWith (flatten cMat) $ \cMatPtr ->
          VS.unsafeWith rhcVec $ \rhcVecPtr ->
            f . g cMatPtr rhcVecPtr $ fromIntegral $ rows cMat

    withWeights f g = maybeWithArray mWeights $ f . g

maybeWithArray :: (Storable a) => Maybe (Vector a) -> (Ptr a -> IO β) -> IO β
maybeWithArray Nothing  f = f nullPtr
maybeWithArray (Just v) f = VS.unsafeWith v f

#if __GLASGOW_HASKELL__ >= 605
{-# INLINE fastMallocForeignPtrArray #-}
fastMallocForeignPtrArray :: forall a. Storable a => Int -> IO (ForeignPtr a)
fastMallocForeignPtrArray n = mallocPlainForeignPtrBytes
                                (n * sizeOf (undefined :: a))
#else
fastMallocForeignPtrArray :: forall a. Storable a => Int -> IO (ForeignPtr a)
fastMallocForeignPtrArray = mallocForeignPtrArray
#endif

--------------------------------------------------------------------------------
-- Minimization options.
--------------------------------------------------------------------------------

-- | Minimization options
data Options r =
    Opts { optScaleInitMu      :: !r -- ^ Scale factor for initial /mu/.
         , optStopNormInfJacTe :: !r -- ^ Stopping thresholds for @||J^T e||_inf@.
         , optStopNorm2Dp      :: !r -- ^ Stopping thresholds for @||Dp||_2@.
         , optStopNorm2E       :: !r -- ^ Stopping thresholds for @||e||_2@.
         , optDelta            :: !r -- ^ Step used in the difference
                                     -- approximation to the Jacobian. If
                                     -- @optDelta<0@, the Jacobian is approximated
                                     -- with central differences which are more
                                     -- accurate (but slower!)  compared to the
                                     -- forward differences employed by default.
         } deriving (Eq, Ord, Read, Show, Data, Typeable)

-- | Default minimization options
defaultOpts :: Fractional r => Options r
defaultOpts = Opts { optScaleInitMu      = c'LM_INIT_MU
                   , optStopNormInfJacTe = c'LM_STOP_THRESH
                   , optStopNorm2Dp      = c'LM_STOP_THRESH
                   , optStopNorm2E       = c'LM_STOP_THRESH
                   , optDelta            = c'LM_DIFF_DELTA
                   }

optsToList :: Options r -> [r]
optsToList (Opts mu  eps1  eps2  eps3  delta) =
                [mu, eps1, eps2, eps3, delta]

--------------------------------------------------------------------------------
-- Constraints
--------------------------------------------------------------------------------

-- | Ensure that these vectors have the same length as the number of parameters.
data Constraints r = Constraints
    { lowerBounds       :: !(Maybe (Params r))            -- ^ Optional lower bounds
    , upperBounds       :: !(Maybe (Params r))            -- ^ Optional upper bounds
    , weights           :: !(Maybe (Params r))            -- ^ Optional weights
    , linearConstraints :: !(Maybe (LinearConstraints r)) -- ^ Optional linear constraints
    } deriving (Read, Show, Typeable)

deriving instance (Eq r, Container Vector r, Num r) => Eq (Constraints r)

-- | Linear constraints consisting of a constraints matrix, @k><m@ and
--   a right hand constraints vector, of length @k@ where @m@ is the number of
--   parameters and @k@ is the number of constraints.
type LinearConstraints r = (Matrix r, Vector r)

-- | * 'mempty' is defined as a 'Constraints' where all fields are 'Nothing'.
--
--   * 'mappend' merges two 'Constraints' by taking the first non-'Nothing' value
--     for each field.
instance Monoid (Constraints r) where
    mempty = Constraints Nothing Nothing Nothing Nothing
    mappend (Constraints lb1 ub1 w1 l1)
            (Constraints lb2 ub2 w2 l2) = Constraints (lb1 `mplus` lb2)
                                                      (ub1 `mplus` ub2)
                                                      (w1  `mplus` w2)
                                                      (l1  `mplus` l2)

--------------------------------------------------------------------------------
-- Output
--------------------------------------------------------------------------------

-- | Information regarding the minimization.
data Info r = Info
  { infNorm2initE      :: !r          -- ^ @||e||_2@             at initial parameters.
  , infNorm2E          :: !r          -- ^ @||e||_2@             at estimated parameters.
  , infNormInfJacTe    :: !r          -- ^ @||J^T e||_inf@       at estimated parameters.
  , infNorm2Dp         :: !r          -- ^ @||Dp||_2@            at estimated parameters.
  , infMuDivMax        :: !r          -- ^ @\mu/max[J^T J]_ii ]@ at estimated parameters.
  , infNumIter         :: !Int        -- ^ Number of iterations.
  , infStopReason      :: !StopReason -- ^ Reason for terminating.
  , infNumFuncEvals    :: !Int        -- ^ Number of function evaluations.
  , infNumJacobEvals   :: !Int        -- ^ Number of jacobian evaluations.
  , infNumLinSysSolved :: !Int        -- ^ Number of linear systems solved,
                                      --   i.e. attempts for reducing error.
  } deriving (Eq, Ord, Read, Show, Data, Typeable)

listToInfo :: (RealFrac r) => [r] -> Info r
listToInfo [a,b,c,d,e,f,g,h,i,j] =
    Info { infNorm2initE      = a
         , infNorm2E          = b
         , infNormInfJacTe    = c
         , infNorm2Dp         = d
         , infMuDivMax        = e
         , infNumIter         = floor f
         , infStopReason      = toEnum $ floor g - 1
         , infNumFuncEvals    = floor h
         , infNumJacobEvals   = floor i
         , infNumLinSysSolved = floor j
         }
listToInfo _ = error "liftToInfo: wrong list length"

-- | Reason for terminating.
data StopReason
  = SmallGradient  -- ^ Stopped because of small gradient @J^T e@.
  | SmallDp        -- ^ Stopped because of small Dp.
  | MaxIterations  -- ^ Stopped because maximum iterations was reached.
  | SingularMatrix -- ^ Stopped because of singular matrix. Restart from current
                   --   estimated parameters with increased 'optScaleInitMu'.
  | SmallestError  -- ^ Stopped because no further error reduction is
                   --   possible. Restart with increased 'optScaleInitMu'.
  | SmallNorm2E    -- ^ Stopped because of small @||e||_2@.
  | InvalidValues  -- ^ Stopped because model function returned invalid values
                   --   (i.e. NaN or Inf). This is a user error.
    deriving (Eq, Ord, Read, Show, Data, Typeable, Enum)

--------------------------------------------------------------------------------
-- Error
--------------------------------------------------------------------------------

data LevMarError
    = LevMarError                    -- ^ Generic error (not one of the others)
    | LapackError                    -- ^ A call to a lapack subroutine failed
                                     --   in the underlying C levmar library.
    | FailedBoxCheck                 -- ^ At least one lower bound exceeds the
                                     --   upper one.
    | MemoryAllocationFailure        -- ^ A call to @malloc@ failed in the
                                     --   underlying C levmar library.
    | ConstraintMatrixRowsGtCols     -- ^ The matrix of constraints cannot have
                                     --   more rows than columns.
    | ConstraintMatrixNotFullRowRank -- ^ Constraints matrix is not of full row
                                     --   rank.
    | TooFewMeasurements             -- ^ Cannot solve a problem with fewer
                                     --   measurements than unknowns.  In case
                                     --   linear constraints are provided, this
                                     --   error is also returned when the number
                                     --   of measurements is smaller than the
                                     --   number of unknowns minus the number of
                                     --   equality constraints.
      deriving (Eq, Ord, Read, Show, Data, Typeable)

-- Handy in case you want to thow a LevMarError as an exception:
instance Exception LevMarError

levmarCErrorToLevMarError :: [(CInt, LevMarError)]
levmarCErrorToLevMarError =
    [ (c'LM_ERROR,                                     LevMarError)
    , (c'LM_ERROR_LAPACK_ERROR,                        LapackError)
  --, (c'LM_ERROR_NO_JACOBIAN,                         can never happen)
  --, (c'LM_ERROR_NO_BOX_CONSTRAINTS,                  can never happen)
    , (c'LM_ERROR_FAILED_BOX_CHECK,                    FailedBoxCheck)
    , (c'LM_ERROR_MEMORY_ALLOCATION_FAILURE,           MemoryAllocationFailure)
    , (c'LM_ERROR_CONSTRAINT_MATRIX_ROWS_GT_COLS,      ConstraintMatrixRowsGtCols)
    , (c'LM_ERROR_CONSTRAINT_MATRIX_NOT_FULL_ROW_RANK, ConstraintMatrixNotFullRowRank)
    , (c'LM_ERROR_TOO_FEW_MEASUREMENTS,                TooFewMeasurements)
  --, (c'LM_ERROR_SINGULAR_MATRIX,                     we don't treat this as an error)
  --, (c'LM_ERROR_SUM_OF_SQUARES_NOT_FINITE,           we don't treat this as an error)
    ]

convertLevMarError :: CInt -> LevMarError
convertLevMarError err = fromMaybe (error $ "Unknown levmar error: " ++ show err)
                                   (lookup err levmarCErrorToLevMarError)


{-# LANGUAGE BangPatterns #-}

module Main where
--import qualified Data.Vector.Storable as V
import Foreign
import Data.Ratio
import Data.Complex

complexZI :: Complex Int
complexZI = 1 :+ 1

ratio23 :: Ratio Int
ratio23 =  1 % 1

putter :: Storable a => a -> Ptr a -> IO a
putter v !ptr = do poke ptr v ; peek ptr

main =
  do
      !vComplex <- alloca (putter complexZI)
      !vRatio <- alloca (putter ratio23)
      if vComplex  == complexZI &&  vRatio  == ratio23
        then putStrLn "success"
        else putStrLn "uh oh, something is wrong with storable"


-- General utilities for working with neural networks.
--
-- Original author: David Banas <capn.freako@gmail.com>
-- Original date:   January 22, 2018
--
-- Copyright (c) 2018 David Banas; all rights reserved World wide.

{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -Wno-unused-top-binds #-}

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeSynonymInstances #-}

{-|
Module      : Haskell_ML.Util
Description : Provides certain general purpose utilities in the Haskell_ML package.
Copyright   : (c) David Banas, 2018
License     : BSD-3
Maintainer  : capn.freako@gmail.com
Stability   : experimental
Portability : ?
-}
module Utils.MLUtils where

import Prelude hiding (zipWith, zip)

import GHC.Generics (Par1(..),(:*:)(..),(:.:)(..))

import           Data.Finite
import           Data.Foldable
import           Data.Function                       (on)
import           Data.Key
import           Data.List                    hiding (zipWith, zip)
import           Data.Ord                            (comparing)
import           GHC.TypeLits
import qualified Data.Vector.Sized            as VS
import qualified Numeric.LinearAlgebra.Data   as LAD
import           Numeric.LinearAlgebra.Static hiding (mean)
import           Text.Printf

import           ConCat.Deep
import           ConCat.Misc                         (sqr, result, cond)
import           ConCat.Orphans                      (fstF, sndF)
import Utils.Utils

-- | Split a list of samples into training/testing sets.
--
-- The Finite given should be the percentage of samples desired for
-- training.
splitTrnTst :: Finite 101 -> [a] -> ([a],[a])
splitTrnTst n xs = splitAt (length xs * (fromInteger . getFinite) n `div` 100) xs

-- | Convert vector of Doubles from sized vector to hmatrix format.
toR :: KnownNat n => V n Double -> R n
toR = vector . VS.toList

-- | Calculate the classification accuracy, given:
--
--   - a functor of results vectors, and
--   - a functor of reference vectors.
classificationAccuracy :: (Foldable f, Keyed f, Eq (Key f), Foldable g, Zip g, Ord a)
                       => g (f a)  -- ^ Functor of result vectors.
                       -> g (f a)  -- ^ Functor of reference vectors.
                       -> Double
classificationAccuracy = (result.result) boolScore (zipWith ((==) `on` maxIndex))
 where boolScore = mean . fmap (cond 1 0)

-- | Mean value of a collection
mean :: (Foldable f, Fractional a) => f a -> a
mean = uncurry (/) . foldr (\e (s,c) -> (e+s,c+1)) (0,0)

-- | Prompt the user for input and read it.
prompt :: Read a => String -> IO a
prompt s = do
  putStrLn s
  readLn

-- | Pretty printer for values of type `R` n.
printVector :: (KnownNat n) => R n -> String
printVector v = foldl' (\ s x -> s ++ printf "%+6.4f  " x) "[ " ((LAD.toList . extract) v) ++ " ]"

-- | Pretty printer for values of type (`R` `m`, `R` `n`).
printVecPair :: (KnownNat m, KnownNat n) => (R m, R n) -> String
printVecPair (u, v) = "( " ++ printVector u ++ ", " ++ printVector v ++ " )"

-- | Show the learning evolution for the various parts/layers in a
-- list of HasLayers.
showPart :: (HasLayers p, RealFrac s, Ord s, PrintfArg s)
         => (p s -> [[s]])  -- ^ The "getter", which grabs the correct part of the parameters.
         -> String          -- ^ A label, which identifies the target of the getter.
         -> [p s]           -- ^ List of parameters, one per epoch.
         -> [String]
showPart f s ps =
  let xss = zip [1::Int,2..] $ (transpose . map f) ps
   in concatMap ( \ (i, xs) ->
                  [ "Average variance in layer " ++ show i ++ " " ++ s ++ ":"
                  , asciiPlot $ map (mean . map sqr) xs
                  ]
                ) xss

-- | Plot a list of Fractionals to an ASCII terminal.
asciiPlot :: (RealFrac a, Ord a, PrintfArg a) => [a] -> String
asciiPlot ys = unlines $
  zipWith (++)
    [ "        "
    , printf " %6.1e " x_max
    , "        "
    , "        "
    , "        "
    , "        "
    , "        "
    , "        "
    , "        "
    , "        "
    , "        "
    , printf " %6.1e " x_min
    , "        "
    ] $
    (:) "^" $ transpose (
    (:) "|||||||||||" $
    for xs $ \x ->
      valToStr $ (x - x_min) * 10 / x_range
    ) ++ ["|0" ++ concat [replicate 16 '_' ++ printf "%4d" (n * length ys `div` 3) | n <- [1..3]] ++ ">"]
      where valToStr  :: (RealFrac a, Ord a) => a -> String
            valToStr x = let i = round (10 - x)
                          in replicate i ' ' ++ "*" ++ replicate (10 - i) ' '
            x_min      = minimum xs
            x_max      = maximum xs
            x_range    = x_max - x_min
            xs         = takeEvery (length ys `div` 60) ys

-- | Return every Nth element of a list.
takeEvery :: Int -> [a] -> [a]
takeEvery n xs = case drop (n-1) xs of
                   (y:ys) -> y : takeEvery n ys
                   []     -> []

-- | Convenience function (= flip map).
for :: Functor f => f a -> (a -> b) -> f b
for = flip fmap

-- | Find the index of the maximum value in a keyed functor.
maxIndex :: (Foldable f, Keyed f, Ord a) => f a -> Key f
maxIndex = fst . maximumBy (comparing snd) . keyed

{--------------------------------------------------------------------
    Network structure
--------------------------------------------------------------------}

-- | A class of parameter types that can be split into a sequence of
-- layers of "weights" and "biases". The meaning of the two notions:
-- "weight" and "bias", are type specific. The only unifyng feature
-- of these, across types, is that they be fully resolved types.
class HasLayers p where
  getWeights :: p s -> [[s]]
  getBiases  :: p s -> [[s]]

instance (HasLayers f, HasLayers g) => HasLayers (g :*: f) where
  getWeights (g :*: f) = getWeights f ++ getWeights g
  getBiases  (g :*: f) = getBiases  f ++ getBiases  g

instance (Foldable a, Foldable b) => HasLayers (a --+ b) where
  getWeights (Comp1 gf) = [(concatMap (toList          . fstF) . toList) gf]
  getBiases  (Comp1 gf) = [(concatMap ((: []) . unPar1 . sndF) . toList) gf]

-- instance (HasLayers f, Foldable g) => HasLayers (g :.: f) where
--   getWeights (Comp1 g) = reverse $ foldl' (\ws -> (: ws) . concat . getWeights) [] g
--   getBiases  (Comp1 g) = reverse $ foldl' (\ws -> (: ws) . concat . getBiases)  [] g

-- instance Foldable f => HasLayers (f :*: Par1) where
--   getWeights (f :*: Par1 _) = [toList f]
--   getBiases  (_ :*: Par1 x) = [[x]]

-- instance Foldable f => HasLayers f where
--   getWeights f = [toList f]
--   getBiases  f = [[]]