{-# LANGUAGE BangPatterns #-}
module STC where

import           Control.Monad.Parallel      as MP
import           Data.Array.IArray           as IA
import           Data.Array.Repa             as R
import           Data.Complex
import           Data.List                   as L
import           Data.Vector.Storable        as VS
import           Data.Vector.Unboxed         as VU
import           FokkerPlanck.FourierSeries
import           FokkerPlanck.GreensFunction
import           FokkerPlanck.Histogram
import           Image.IO
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Parallel
import           Utils.Time

main = do
  args@(deviceIDsStr:numPointStr:deltaStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:deltaLogStr:taoStr:numR2FreqStr:periodR2Str:phiFreqsStr:rhoFreqsStr:thetaFreqsStr:scaleFreqsStr:initDistStr:initScaleStr:histFilePath:stdStr:numBatchR2FreqsStr:numBatchOriStr:numThreadStr:_) <-
    getArgs
  let deviceIDs = read deviceIDsStr :: [Int]
      numPoint = read numPointStr :: Int
      delta = read deltaStr :: Double
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      tao = read taoStr :: Double
      numR2Freq = read numR2FreqStr :: Int
      periodR2 = read periodR2Str :: Double
      phiFreq = read phiFreqsStr :: Int
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreq = read rhoFreqsStr :: Int
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      thetaFreq = read thetaFreqsStr :: Int
      thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Int
      scaleFreqs = L.map fromIntegral [-scaleFreq .. scaleFreq]
      initScale = read initScaleStr :: Double
      initDist = read initDistStr :: [(Double, Double, Double, Double)]
      initPoints = L.map (\(x, y, t, s) -> Point x y t s) initDist
      numThread = read numThreadStr :: Int
      folderPath = "output/test/FourierSeriesOfPinwheels"
      maxScale = read maxScaleStr :: Double
      halfLogPeriod = log maxScale
      deltaLog = read deltaLogStr :: Double
      std = read stdStr :: Double
      numBatchR2Freqs = read numBatchR2FreqsStr :: Int
      numBatchOri = read numBatchOriStr :: Int
      radius = fromIntegral $ div numPoint 2
  removePathForcibly folderPath
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  printCurrentTime "Start computing coefficients..."
  hist <-
    computeFourierCoefficients
      deviceIDs
      numPoints
      numOrientation
      delta
      initScale
      thetaSigma
      tao
      numR2Freq
      periodR2
      std
      numBatchR2Freqs
      numBatchOri
      phiFreq
      rhoFreq
      thetaFreq
      scaleFreq
  hist <-
    sampleCartesian
      folderPath
      histFilePath
      deviceIDs
      radius
      144
      0.25 -- deltaLog
      initScale
      thetaSigma
      tao
      phiFreqs
      rhoFreqs
      thetaFreqs
      scaleFreqs
  printCurrentTime "Done"
  printCurrentTime "Start Convloution.."
  let !coefficients =
        normalizeFreqArr' std phiFreqs rhoFreqs . getNormalizedHistogramArr $
        hist
      initialDistribution =
        computeInitialDistributionFull' r2Freq radius phiFreq rhoFreq initPoints
      !harmonicsArray' =
        computeFourierSeriesOfLogPolarHarmonicsArray
          radius
          deltaLog
          r2Freq
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
          halfLogPeriod :: IA.Array (Int, Int) (VS.Vector (Complex Double))
      -- !harmonicsArray' = computeHarmonicsArray numPoint 1 numPoint 1 phiFreqs rhoFreqs thetaFreqs scaleFreqs 1 radius
      source = convolveFull' coefficients harmonicsArray' initialDistribution
      inverseR2Harmonics =
        computeRectangularInverseHarmonics numPoint numPoint delta radius r2Freq :: [VS.Vector (Complex Double)]
      -- pinwheelVec =
      --   parMap
      --     rdeepseq
      --     (\(idx, vec) ->
      --        ( idx
      --        , VU.fromList . L.map (VS.sum . VS.zipWith (*) vec) $
      --          inverseR2Harmonics)) .
      --   assocs $
      --   harmonicsArray'
  -- print initialDistribution
  createDirectoryIfMissing True (folderPath </> "Pinwheels")
  plotImageRepaComplex (folderPath </> "Pinwheels" </> "Pinwheel_(1,1).png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoint :. numPoint) .
    VU.fromList . L.map (VS.sum . VS.zipWith (*) (harmonicsArray' IA.! (1, 1))) $
    inverseR2Harmonics
  -- MP.mapM_
  --   (\((a, b), vec) ->
  --      plotImageRepaComplex
  --        (folderPath </> "Pinwheels" </> printf "Pinwheel_(%d,%d).png" a b) .
  --      ImageRepa 8 . fromUnboxed (Z :. (1 :: Int) :. numPoint :. numPoint) $
  --      vec)
  --   pinwheelVec
  -- printCurrentTime "Done"
  printCurrentTime "Start ploting.."
  plotDFTArrayPower
    (folderPath </> "SourcePower.png")
    numPoint
    numPoint
    -- (2 * r2Freq + 1)
    -- (2 * r2Freq + 1)
    source
  plotFourierCoefficientsPower
    (folderPath </> "Source.png")
    numPoint
    numPoint
    inverseR2Harmonics
    source
  printCurrentTime "Done"


{-# language FlexibleInstances #-}
{-|
Module      : MachineLearning.Model.Fitness
Description : TIR expression data structures
Copyright   : (c) Fabricio Olivetti de Franca, 2022
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX

This module exports the functions that calculates the coefficients,
evaluate the fitness, penalty and constraints of a TIR expression.
-}
module MachineLearning.Model.Fitness
  ( evalTrain
  , evalTest
  , selectValidTerms
  ) where

import Data.Bifunctor
import Data.Maybe            (fromJust)
import Data.Vector.Storable  (Vector)
import Data.Vector           ((!))
import Numeric.LinearAlgebra ((<\>))
import Numeric.ModalInterval      (Kaucher, inf, sup, width, singleton)
import qualified Numeric.ModalInterval as Interval
import Control.Monad.Reader
import qualified Data.Vector.Storable  as VS
import qualified Data.Vector           as V
import qualified Numeric.LinearAlgebra as LA
import Data.Maybe (fromMaybe)
import MachineLearning.Model.Measure       (Measure)
import MachineLearning.Model.Regression    (nonlinearFit, evalPenalty, fitTask, predictTask, applyMeasures)
import MachineLearning.TIR       (TIR(..),  Individual(..), Dataset, Constraint, assembleTree, replaceConsts)
import MachineLearning.Utils.Config       (Task(..), Penalty)
import Data.SRTree                (SRTree(..), Function, OptIntPow(..), evalTree, evalTreeMap, evalFun, inverseFunc, countNodes)

import Data.SRTree.Print

-- | removes invalid terms from the TIR expression. Invalid terms
-- are those that evaluate to `NaN` or `Infinite` within the
-- domains of each variable. The domains are either provided by
-- the configuration file or estimated using the training data.
selectValidTerms :: TIR -> V.Vector (Kaucher Double) -> TIR
selectValidTerms tir@(TIR _ p q) domains = tir{ _p=p', _q=q' }
  where
    p' = selectValid p
    q' = selectValid q

    selectValid = filter (\(_, g, ps) -> isValidInterval $ evalFun g (evalPi ps))
    evalPi      = foldr (\(ix, k) acc -> acc * (domains ! ix ^. k)) 1
{-# INLINE selectValidTerms #-}

isValidInterval = not. isInvalidInterval
{-# INLINE isValidInterval #-}

isInvalidInterval :: Kaucher Double -> Bool
isInvalidInterval ys =  Interval.isEmpty ys
                     || Interval.isInvalid ys
                     || isInfinite ys1 || isInfinite ys2
                     || ys2 < ys1
                     || abs ys1 >= 1e50 || abs ys2 >= 1e50
                     || isNaN ys1 || isNaN ys2
                     || width ys < 1e-8
  where
    ys1 = fromMaybe (-1/0) $ inf ys
    ys2 = fromMaybe (1/0) $ sup ys
{-# INLINE isInvalidInterval #-}

-- | evaluates an individual first fitting the expression
-- with either OLS or a nonlinear optimization (not yet implemented)
-- and calculating the fitness vector, constraints, penalty.
evalTrain :: Task                          -- ^ Regression or Classification task
          -> Bool                          -- ^ if we are fitting the final best individual, in this case do not split the training data for validation
          -> [Measure]                     -- ^ list of performance measures to calculate
          -> Constraint                    -- ^ constraint function
          -> Penalty                       -- ^ penalty
          -> V.Vector (Kaucher Double)     -- ^ variable domains represented as a Kaucher Interval
          -> Dataset Double                -- ^ training data
          -> Vector Double                 -- ^ training target
          -> Dataset Double                -- ^ validation data
          -> Vector Double                 -- ^ validation target
          -> Individual
          -> Individual
evalTrain task isRefit measures cnstrFun penalty domains xss_train ys_train xss_val ys_val sol
--  | LA.cols zss == 0                   = error "found"
--  | (not.null) (LA.find (\x -> isNaN x || isInfinite x) zss)  = error $ (show $ _chromo sol) <> show domains
  | not isRefit && (not.null._fit) sol = sol
--  | LA.cols zssP == 0                  = sol { _fit = [1/0] }
  | otherwise                          = sol{ _chromo  = fitted
                                            , _fit     = fitness
                                            , _weights = ws
                                            , _constr  = cnst
                                            , _len     = len
                                            , _penalty = pnlty
                                            }
  where
    -- Fit the rational IT
    tir          = selectValidTerms (_chromo sol) domains
    ws           = fitTask task tir xss_train ys_train

    -- Validate (it should be applied to every different weights set)
    fitted         = replaceConsts tir . V.tail . VS.convert . head
                   $ ws
    fitness        = map nan2inf . fromJust . evalTest task measures xss_val ys_val
                   $ sol{ _chromo=fitted, _weights=ws }
    -- Length and constraint
    tree           = assembleTree (V.head $ VS.convert $ head ws) fitted
    len            = countNodes tree
    cnst           = cnstrFun tree
    pnlty          = evalPenalty penalty len cnst

-- | Evaluates an expression into the test set. This is different from `fitnessReg` since
-- it doesn't apply OLS.
evalTest :: Task -> [Measure] -> Dataset Double -> Vector Double -> Individual -> Maybe [Double]
evalTest task measures xss ys sol
  | null weights = Nothing
  | otherwise    = Just
                 $ applyMeasures measures ys
                 $ predictTask task
                 $ map (evalTIR xss' bias . replaceConsts tir . V.tail . VS.convert) weights
  where
    tir          = _chromo sol
    weights      = _weights sol
    bias         = V.head $ VS.convert $ head weights -- only works for regression
    -- nSamples     = VS.length $ V.head xss
    xss'         = V.tail xss
    -- treeEval t   = fromJust $ runReader (evalTreeMap (VS.replicate nSamples) t) (xss' V.!?)

evalTIR :: Dataset Double -> Double -> TIR -> LA.Vector Double
evalTIR xss bias (TIR g p q) = evalFun g ((LA.scalar bias + p') / (1 + q'))
  where
    p'     = foldr (\(w, h, ks) acc -> LA.scalar w * evalFun h (evalPi ks) + acc) 0 p
    q'     = foldr (\(w, h, ks) acc -> LA.scalar w * evalFun h (evalPi ks) + acc) 0 q
    evalPi = foldr (\(ix,k) acc -> acc * (xss ! ix)^^k) 1
{-# INLINE evalTIR #-}

--instance OptIntPow (LA.Vector Double) where
--  (^.) = (^^)

nan2inf :: RealFloat a => a -> a
nan2inf x | isNaN x = 1/0
          |otherwise = x
{-# INLINE nan2inf #-}


module Data.MyPolynomial.Print
	where

import Data.MyPolynomial.Type
import Data.Complex (realPart, imagPart, Complex)
import Data.List (sort)
import Data.IntMap.Lazy (lookupLT, lookupGE, IntMap)

 -- `procelain' functions are intended for canonical printing :
 -- i.e. Show instances, i.e. serialization, i.e. Read reciprocal

 -- `pretty' functions are intended for pretty printing only :
 -- i.e. absolutly no consistent Read <=> Show reciprocity.

 -- `prints' functions are general purpose printing functions;
 -- they either use `pretty' or `porcelain' to generate output.
 -- The 3rd person-like terminal 's' is a general Haskell convention
 -- for functions using a provided interface to do a specific job.

 -- S suffixed functions return a String difference list (ShowS)
 -- for efficient arbitrary concatenation.

 -- M suffixed functions take a polynomial represented by
 -- an IntMap Float instance.

---------------------------------------------------------------------------
----  Abstract printers
---------------------------------------------------------------------------

printsPolynomialS :: (Monomial -> ShowS) -> Polynomial -> ShowS
printsPolynomialS _ []  = ('0':)
printsPolynomialS pfS (mn:[]) = pfS mn
printsPolynomialS pfS (m1:(m2@(c2 :*^: p2)):ms) = pfS m1 . signBridge . next
 where
  signBridge
   | c2 < 0 = (" - " ++)
   | otherwise = (" + " ++)
  next
   | c2 < 0 = printsPolynomialS pfS (((-c2) :*^: p2):ms)
   | otherwise = printsPolynomialS pfS (m2:ms)

printsPolynomialSM :: (Polynomial -> ShowS) -> IntMap Float -> ShowS
printsPolynomialSM pfS mmap = ( pfS . sort ) $ consUp (smallest 0 mmap) mmap ([]++) []
  where
    smallest k map = case lookupLT k map of { Nothing -> k;	Just (k, _) -> smallest k map; }
    consUp k map dl = case lookupGE k map of	Nothing -> dl
						Just (kn, v) -> consUp (kn + 1) map (((v :*^: kn):) . dl)

printsEquationS :: (Polynomial -> ShowS) -> Equation -> ShowS
printsEquationS pfS (Eq (l, r)) = memb l . showString " = " . memb r
 where
  memb [] = ('0':)
  memb p  = pfS p

---------------------------------------------------------------------------
----  Canonical printing ( Read <=> Show )
---------------------------------------------------------------------------

porcelainMonomialS :: Monomial -> ShowS
porcelainMonomialS (c :*^: p) = shows c . showString " * X^" . shows p

porcelainPolynomialS :: Polynomial -> ShowS
porcelainPolynomialS = printsPolynomialS porcelainMonomialS

porcelainEquationS :: Equation -> ShowS
porcelainEquationS = printsEquationS porcelainPolynomialS

porcelainPolynomialSM :: IntMap Float -> ShowS
porcelainPolynomialSM = printsPolynomialSM porcelainPolynomialS

---------------------------------------------------------------------------
----  Pretty printing
---------------------------------------------------------------------------

prettyMonomialS :: Monomial -> ShowS
prettyMonomialS (0 :*^: _) = showString ""
prettyMonomialS (c :*^: 0) = shows c
prettyMonomialS (c :*^: 1) = shows c . showString " X"
prettyMonomialS (c :*^: p) = prettyMonomialS (c :*^: 1) . ('^':) . shows p

prettyMonomial :: Monomial -> String
prettyMonomial mn =  prettyMonomialS mn ""

prettyPolynomialS :: Polynomial -> ShowS
prettyPolynomialS = ( printsPolynomialS prettyMonomialS ) . purgeHollow
  where purgeHollow = filter ( \(c :*^: _) -> c /= 0 )

prettyPolynomial :: Polynomial -> String
prettyPolynomial p = prettyPolynomialS p ""

prettyPolynomialSM :: IntMap Float -> ShowS
prettyPolynomialSM = printsPolynomialSM prettyPolynomialS

prettyPolynomialM :: IntMap Float -> String
prettyPolynomialM m = prettyPolynomialSM m []

prettyEquationS :: Equation -> ShowS
prettyEquationS = printsEquationS (printsPolynomialS porcelainMonomialS)

prettyEquation :: Equation -> String
prettyEquation eq = prettyEquationS eq []

prettyComplexS :: (Ord c, Num c, Show c, RealFloat c) => Complex c -> ShowS
prettyComplexS c = shows a . img
	where
	 a = realPart c
	 img
	  | b < 0 = (" - " ++) . shows (-b) . ('i':)
	  | b > 0 = (" + " ++) . shows b . ('i':)
	  | b == 0 = showString ""
	  where
	   b = imagPart c

prettyComplex :: (Ord c, Num c, Show c, RealFloat c) => Complex c -> String
prettyComplex c = prettyComplexS c ""

---------------------------------------------------------------------------
----  Show instances
---------------------------------------------------------------------------

instance Show Monomial where
  showsPrec _ = porcelainMonomialS
  showList = porcelainPolynomialS

instance Show Equation where
  showsPrec _ = porcelainEquationS


module HVX.Internal.Solvers
  ( constStep
  , decNonSumStep
  , subgradMinimize
  , subgradMaximize
  , subgradLoop
  , ellipsoidMinimize
  , ellipsoidMaximize
  , ellipsoidLoop
  ) where

import Data.List
import Data.Ord
import Data.Maybe
import Numeric.LinearAlgebra

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

class CanMinimize vex
instance CanMinimize Affine
instance CanMinimize Convex

class CanMaximize vex
instance CanMaximize Affine
instance CanMaximize Concave

-- | Constant step size.
constStep :: Double -> Int -> Double
constStep alpha _ = alpha

-- | 1/n step size.
decNonSumStep :: Double -> Int -> Double
decNonSumStep alpha itr = alpha / fromIntegral itr

-- | Use subgradient method to solve a minimization problem.
subgradMinimize :: CanMinimize vex =>
     Expr vex mon     -- ^ Objective to minimize.
  -> [Constraint]     -- ^ Constraints on minimization problem.
  -> (Int -> Double)  -- ^ Stepsize function.
  -> Int              -- ^ Number of iterations to run.
  -> Vars             -- ^ Variables and their starting values.
  -> (Vars, Double)   -- ^ Optimal variable values and optimal objective value.
subgradMinimize = subgradLoop 1

-- | Use subgradient method to solve a maximization problem.
subgradMaximize :: CanMaximize vex =>
     Expr vex mon     -- ^ Objective to maximize.
  -> [Constraint]     -- ^ Constraints on maximization problem.
  -> (Int -> Double)  -- ^ Stepsize function.
  -> Int              -- ^ Number of iterations to run.
  -> Vars             -- ^ Variables and their starting values.
  -> (Vars, Double)   -- ^ Optimal variable values and optimal objective value.
subgradMaximize objective constraints stepFun maxItr vars = (fst result, negate $ snd result)
  where result = subgradLoop 1 (EFun Neg objective) constraints stepFun maxItr vars

-- | Subgradient solver that minimizes an expression subject to constraints.
subgradLoop :: Int -> Expr vex mon -> [Constraint] -> (Int -> Double) -> Int -> Vars
  -> (Vars, Double)
subgradLoop itr objective constraints stepFun maxItr vars =
  if itr >= maxItr || vars == varsNext
    then (vars, evaluate objective vars @@> (0,0))
    else subgradLoop (itr + 1) objective constraints stepFun maxItr varsNext
      where varsNext = map (updateWithSubgrad objective constraints (stepFun itr) vars) vars

updateWithSubgrad :: Expr vex mon -> [Constraint] -> Double -> Vars -> (Var, Mat) -> (Var, Mat)
updateWithSubgrad objective constraints stepSize vars (varname, val) =
  if any isNaN (toList $ flatten g)
    then (varname, val)
    else (varname, val - scale stepSize g)
      where g = fromMaybe (trans $ jacobianWrtVar objective vars varname) $
                          getConstraintSubgrad constraints vars varname

type Soid = (Var, Mat, Mat)

-- | Use ellipsoid method to solve a minimization problem.
ellipsoidMinimize :: CanMinimize vex =>
     Expr vex mon            -- ^ Objective to minimize.
  -> [Constraint]            -- ^ Constraints on minimization problem.
  -> [(Var,Int)]             -- ^ Variable names and their sizes.
  -> Double                  -- ^ Tolerance.
  -> Double                  -- ^ Radius to seed initial sphere with.
  -> (Vars, Double, Double)  -- ^  Optimal variable values, lower bound, and upper bound.
ellipsoidMinimize objective constraints varsizes tol radius =
  ellipsoidLoop objective constraints tol soids (-infinity) infinity
    where soids = map (\(name,n) -> (name, zeroVec n, scale radius $ ident n)) varsizes

-- | Use ellipsoid method to solve a maximization problem.
ellipsoidMaximize :: CanMaximize vex =>
     Expr vex mon            -- ^ Objective to maximize.
  -> [Constraint]            -- ^ Constraints on maximization problem.
  -> [(Var,Int)]             -- ^ Variable names and their sizes.
  -> Double                  -- ^ Tolerance.
  -> Double                  -- ^ Radius to seed initial sphere with.
  -> (Vars, Double, Double)  -- ^  Optimal variable values, lower bound, and upper bound.
ellipsoidMaximize objective constraints varsizes tol radius = (optvars, negate ubound, negate lbound)
  where (optvars, lbound, ubound) = ellipsoidLoop (EFun Neg objective) constraints tol soids (-infinity) infinity
        soids = map (\(name,n) -> (name, zeroVec n, scale radius $ ident n)) varsizes

-- | Ellipsoid method solver that minimizes an expression subject to constraints.
ellipsoidLoop :: Expr vex mon -> [Constraint] -> Double -> [Soid] -> Double -> Double
  -> (Vars, Double, Double)
ellipsoidLoop objective constraints tol soids lbound ubound =
  if ubound - lbound < tol || sqrtgpgsum == 0.0
    then (centers, nlbound, nubound)
    else ellipsoidLoop objective constraints tol nsoids nlbound nubound
      where updates = map (updateEllipsoid objective constraints centers) soids
            nsoids = map fst updates
            sqrtgpgsum = sqrt . sum . map snd $ updates
            objval = evaluate objective centers @@> (0,0)
            nlbound = max lbound $ objval - sqrtgpgsum
            nubound = min ubound objval
            centers = map (\(name, val, _) -> (name, val)) soids

updateEllipsoid :: Expr vex mon -> [Constraint] -> Vars -> Soid -> (Soid, Double)
updateEllipsoid objective constraints vars soid =
  if gpg == 0.0
    then (soid, objgpg)
    else (nsoid, objgpg)
      where objg = trans $ jacobianWrtVar objective vars varname
            g = fromMaybe objg $ getConstraintSubgrad constraints vars varname
            n = fromIntegral $ rows val
            gpg = (@@> (0,0)) $ trans g <> p <> g
            objgpg = (@@> (0,0)) $ trans objg <> p <> objg
            gtilde = scale (1 / sqrt gpg) g
            nval = val - scale (1 / (n + 1)) (p <> gtilde)
            np = scale (n**2 / (n**2 - 1)) $
              p - scale (2/(n+1)) (p <> (gtilde <> trans gtilde) <> p)
            (varname, val, p) = soid
            nsoid = (varname, nval, np)

-- | Get the subgradient of the maximally violated constraint if there is a
-- violation. Returns nothing if all constraints are satisfied
getConstraintSubgrad :: [Expr vex mon] -> Vars -> String -> Maybe Mat
getConstraintSubgrad constraints vars var =
  if null constraints || maxVal < 0.0
    then Nothing
    else Just cg
      where cg = trans $ jacobianWrtVar maxConstraint vars var
            (maxConstraint, maxVal) = maximumBy (comparing snd) constraintViolations
            constraintViolations = map (\c -> (c, evaluate c vars @@> (0,0))) constraints


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
module Statistics.Distribution.Random.Exponential (
    exponential
) where

import Random.CRI
import Statistics.Distribution.Random.Uniform

-- q[k-1] = sum(log(2)^k / k!)  k=1,..,n,
q :: [Double]
-- q = let factorial = foldr (*) 1 . enumFromTo 1
--         qk :: Integer -> Double
--         qk k = (log 2)^k / (fromIntegral (factorial k))
--         qs :: [Double]
--         qs = map qk [1..]
--     in map (\ n -> sum (take n qs)) (take 16 [1..])
q = [0.6931471805599453,
     0.9333736875190459,
     0.9888777961838675,
     0.9984959252914960,
     0.9998292811061389,
     0.9999833164100727,
     0.9999985691438767,
     0.9999998906925558,
     0.9999999924734159,
     0.9999999995283275,
     0.9999999999728814,
     0.9999999999985598,
     0.9999999999999289,
     0.9999999999999968,
     0.9999999999999999,
     1.0000000000000000]
{-# INLINE q #-}

exponential :: Source m g Double => g m -> m Double
exponential rng =
    let !q0 = q !! 0
        mk_au :: Double -> Double -> (Double, Double)
        mk_au !a !u
            | u > 1.0   = (a, u - 1.0)
            | otherwise = mk_au (a + q0) (u + u)

        go a _ umin ![] = return (a + umin * q0)
        go a u umin !(qi:qs) = do
            !ustar <- uniform rng
            let umin' = min ustar umin
            if u > qi
                then go a u umin' qs
                else return (a + umin' * q0)
    in do
        !u' <- uniform rng
        let !(a, u) = mk_au 0.0 (u' + u')
        if u <= q0
            then return (a + u)
            else do
                !us <- uniform rng
                go a u us (tail q)


{-# LANGUAGE UnicodeSyntax, DataKinds, TypeOperators, KindSignatures,
             TypeInType, GADTs, MultiParamTypeClasses, FunctionalDependencies,
             TypeFamilies, AllowAmbiguousTypes, FlexibleInstances,
             UndecidableInstances, InstanceSigs, TypeApplications,
             ScopedTypeVariables, EmptyCase, FlexibleContexts,
             RankNTypes, LambdaCase
#-}

-- Implementation of density matrix interpretation of quantum computation
module LinTrans where

import Prelim

import qualified Data.Complex as C
import Data.Complex (Complex(..))
import Control.Monad.State.Lazy
import Data.Singletons
import Data.Maybe (fromJust)
--import Data.Tuple hiding (swap)
import Unsafe.Coerce
import Data.Constraint (Dict(..))
import Data.List ( (\\) )
import qualified Debug.Trace as Tr

-- I don't like the show instance for ℂ
newtype ℂ = ℂ (Complex Double)

liftC :: (Complex Double -> Complex Double) -> ℂ -> ℂ
liftC f (ℂ c) = ℂ $ f c

i :: ℂ
i = ℂ $ 0 :+ 1
conjugate = liftC C.conjugate
instance Num ℂ where
  ℂ m + ℂ n = ℂ $ m + n
  ℂ m - ℂ n = ℂ $ m - n
  ℂ m * ℂ n = ℂ $ m * n
  abs = liftC abs
  signum = liftC signum
  fromInteger = ℂ . fromInteger
instance Fractional ℂ where
  fromRational = ℂ . fromRational
  ℂ m / ℂ n    = ℂ $ m / n
instance Floating ℂ where
  pi    = ℂ $ pi
  exp   = liftC exp
  log   = liftC log
  sin   = liftC sin
  cos   = liftC cos
  asin  = liftC asin
  acos  = liftC acos
  atan  = liftC atan
  sinh  = liftC sinh
  cosh  = liftC cosh
  asinh = liftC asinh
  acosh = liftC acosh
  atanh = liftC atanh
instance Show ℂ where
  show (ℂ (α :+ β)) =   if β == 0 then show α
                        else if α == 0 then show β ++ "i"
                        else show α ++ " + " ++ show β ++ "i"

-- A matrix of dimension m × n is a function from
-- pairs of nats bounded by m and n respectively.
type Matrix m n = (BNat m,BNat n) -> ℂ
type Squared n = Matrix (Two `RaiseTo` n) (Two `RaiseTo` n)

matrix :: forall m n. (SingI m, SingI n) => [ℂ] -> Matrix m n
matrix ls (i,j) = ls !! (n * fromIntegral i + fromIntegral j)
  where
    n = toInt (sing :: Sing n)

--------------
-- Printing --
--------------
rows :: forall m n. (SingI m, SingI n) => Matrix m n -> [[ℂ]]
rows mat = fmap f $ allBNat (sing :: Sing m)
  where
    f i = fmap (\j -> mat (i,j)) $ allBNat (sing :: Sing n)

showRows :: (SingI m, SingI n) => Matrix m n -> [String]
showRows mat = show <$> rows mat

instance (SingI m, SingI n) => Show (Matrix m n) where
  show mat = unlines $ showRows mat

---------------------------------------
-- Primitive matrices and operatoins --
---------------------------------------

ident :: Matrix m m
ident (i,j) = if i==j then 1 else 0

transpose :: Matrix m n -> Matrix n m
transpose mat (i,j) = conjugate $ mat (j,i)

{-
-- if i : BNat m then raise m n = i : BNat (S m)
-- so pruneRows mat is just the first m rows of mat
pruneRows :: Sing m -> Matrix (S m) n -> Matrix m n
pruneRows m mat (i,j) = mat (raise m i,j)

pruneCols :: Sing n -> Matrix m (S n) -> Matrix m n
pruneCols n mat (i,j) = mat (i,raise n j)
-}

sumTo :: Sing n -> BNat n -> (BNat n -> ℂ) -> ℂ
sumTo _ (BNat SZ)     _ = 0
sumTo n (BNat (SS m)) f = f m' + sumTo n m' f
  where
    m' = case succLtTrans m n of Dict -> BNat m -- m < S m < n

dot :: Sing n -> Matrix (S Z) n -> Matrix n (S Z) -> ℂ
dot n v1 v2 = sumTo n (maxBNat n) (\x -> v1 (0,x) * v2(x,0))

row :: BNat m -> Matrix m n -> Matrix (S Z) n
row i mat (_,j) = mat (i,j)

col :: BNat n -> Matrix m n -> Matrix m (S Z)
col j mat (i,_) = mat (i,j)

plus :: Matrix m n -> Matrix m n -> Matrix m n
plus mat1 mat2 (i,j) = mat1(i,j) + mat2(i,j)

mult :: forall m n p. SingI n => Matrix m n -> Matrix n p -> Matrix m p
mult mat1 mat2 (i,j) =
--  Tr.trace ("Composing for intermediate " ++ show (sing :: Sing n)) $
    dot sing (row i mat1) (col j mat2)

kron :: forall m1 m2 n1 n2. (SingI m1, SingI n1, SingI m2, SingI n2)
     => Matrix m1 n1 -> Matrix m2 n2 -> Matrix (m1 `Times` m2) (n1 `Times` n2)
kron mat1 mat2 (BNat i,BNat j) =
--    Tr.trace ("Computing ⊗ for " ++ show (m1, n1) ++ "×" ++ show (m2, n2)) $
    mat1 (bNat i1, bNat j1) * mat2 (bNat i2,bNat j2)
  where
    i1 = i `divSNat` m2
    i2 = i `modSNat` m2
    j1 = j `divSNat` n2
    j2 = j `modSNat` n2
    m2 = (sing :: Sing m2)
    n2 = (sing :: Sing n2)

trace :: forall m. SingI m => Matrix m m -> ℂ
trace mat = foldr f 0 ls
  where
    ls = allBNat (sing :: Sing m)
    f i x = mat (i,i) + x

----------------------
-- Quantum Matrices --
----------------------

square :: forall n. SingI n => [ℂ] -> Squared n
square ls = withSingI (two `raiseToSNat` (sing :: Sing n)) $ matrix ls

density0 :: Matrix Two Two
density0 = matrix [1,0
                  ,0,0]

density1 :: Matrix Two Two
density1 = matrix [0,0
                  ,0,1]

hadamard = square @One $ [1/sqrt 2, 1/sqrt 2
                         ,1/sqrt 2, -1/sqrt 2]

pauliX = square @One $ [0,1,1,0]
pauliY = square @One $ [0,-i,i,0]
pauliZ = square @One $ [1,0,0,-1]

cnotD  = square @Two $ [1,0,0,0
                       ,0,1,0,0
                       ,0,0,0,1
                       ,0,0,1,0]

newD :: Bool -> Squared One
newD True  = density1
newD False = density0

--------------------------------------------------------------------

----------------------
-- Density Matrices --
----------------------

data Density where
  Density :: Sing n -> Squared n -> Density
instance Show Density where
  show (Density n m) = withSingI (two `raiseToSNat` n) $ show m

size :: forall m. SingI m => Matrix m m -> Int
size _ = toInt (sing :: Sing m)
logsizeD :: Density -> Int
logsizeD (Density n mat) = toInt n
sizeD :: Density -> Int
sizeD ρ = 2 ^ logsizeD ρ

transposeD :: Density -> Density
transposeD (Density m ρ) = Density m $ transpose ρ

multD :: Density -> Density -> Density
multD = undefined

kronD :: Density -> Density -> Density
kronD (Density m mat1) (Density n mat2) =
      withSingI (raiseToSNat two n) $
      withSingI (raiseToSNat two m) $
      withSingI (raiseToSNat two (plusSNat m n)) $
      case raiseToPlus two m n of {Dict ->
        Density (m `plusSNat` n) $ mat1 `kron` mat2
      }

identD :: Int -> Density
identD n = case fromIntegral @_ @(SomeSing Nat) n of SomeSing n' -> Density n' ident

nil :: Int -> Density
nil n = case fromIntegral n of
          SomeSing n' -> withSingI (two `raiseToSNat` n') $
                         Density n' . matrix $ fromIntegral <$> repeat 0

------------------
-- Permutations --
------------------

-- encode and decode modulo a key k
encode :: Int -> [Int] -> Int
encode k [] = 0
encode k (i : ls) = i + k * encode k ls

decode :: Int -> Int -> [Int]
decode k i | i == 0 = []
decode k i | otherwise = i `mod` k : decode k (i `div` k)

fromAssocList :: [(Int,Int)] -> Int -> Int
fromAssocList []           = id
fromAssocList ((a,b) : ls) = assocFun a b . fromAssocList ls
  where
    assocFun a b i = if i == a then b else i

-- Check if two lists are equal up to the permutation specified, i.e.
-- if permuting ls1 by f is equal to ls2
isPermutation :: Eq a => (Int -> Int) -> [a] -> [a] -> Bool
isPermutation f ls1 ls2 = all (\i -> ls1 !! i == ls2 !! f i) [0..len]
  where
    len = length ls1

swapFun :: forall n. SingI n => (Int -> Int) -> Squared n
swapFun f (i,j) = if isPermutation f (decode n' i') (decode n' j') then 1 else 0
  where
    n' = toInt (sing :: Sing n)
    i' = toInt i
    j' = toInt j

-- swap takes a list of qubit/bit variables and produces a matrix that permutes
-- those qubits.
-- i.e. swap ls |φ_0 ⋯ φ_n⟩ = |φ_f(0) ⋯ ρ_f(n)⟩
-- where f = fromAssocList ls
swap :: forall n. SingI n => [Int] -> Squared n
swap ls = swapFun @n (fromAssocList $ zip [0..] ls)

-------------------
-- Density Monad --
-------------------

-- A density monad is a nondeterminism state monad:
-- Density -> [Density]
-- An element op of type DensityMonad corresponds to the superoperator
-- \ρ -> ∑ (op ρ)
type DensityMonad = StateT Density []

newM :: Bool -> DensityMonad Int
newM b = do
    ρ <- get
    put $ ρ `kronD` (Density one $ newD b)
    return $ sizeD ρ

-- runQ applies the superoperator to the identity density matrix of size 1.
runQ :: DensityMonad a -> [(a,Density)]
runQ m = runStateT m (Density SZ ident)

-- getDensity combines the result of runQ into a single density matrix
getDensity :: DensityMonad a -> Density
getDensity m = foldr f (nil n) ls
  where
    ls = runQ m
    n = logsizeD . snd $ head ls
    f (_,ρ) ρ0 = ρ + ρ0

-----------------
-- Application --
-----------------

superM :: forall m n. Matrix (Two `RaiseTo` m) (Two `RaiseTo` n) -> DensityMonad ()
superM mat = do Density _ ρ ← get
                put . Density _ $ mat `mult` ρ `mult` transpose mat
  where
    matD = Density (sing :: Sing m) mat

applyUnitaryM :: forall m. SingI m => Squared m -> [Int] -> DensityMonad ()
applyUnitaryM u ls = do superM @m σ -- moves the relevant qubits to the front of the density matrix
                        superM @m u -- applies operation
                        superM @m $ transpose σ -- moves the qubits back
  where
    σ = swap @m ls

applyMatrixM :: Matrix (Two `RaiseTo` m) (Two `RaiseTo` n)
             -> [Int] -> DensityMonad ()
applyMatrixM mat ls = undefined

measM :: Int -> DensityMonad Bool
measM i = do
    ρ <- get
    (b,ρ') <- lift [ (False, applyMatrixD @One density0 [i] ρ)
                   , (True,  applyMatrixD @One density1 [i] ρ) ]
    put ρ'
    return b

------------------
-- Num instance --
------------------

instance (SingI m, SingI n) => Num (Matrix m n) where
  (+) mat1 mat2 (i,j) = mat1(i,j) + mat2(i,j)
  (-) mat1 mat2 (i,j) = mat1(i,j) - mat2(i,j)
  (*) mat1 mat2 (i,j) = mat1(i,j) * mat2(i,j)
  abs mat (i,j) = abs $ mat(i,j)
  signum mat = undefined -- trace mat
  fromInteger n = matrix . repeat $ fromIntegral n
instance Num Density where
  Density n1 mat1 + Density n2 mat2 =
    withSingI (two `raiseToSNat` n1) $
    withSingI (two `raiseToSNat` n2) $
    case eqSNat n1 n2 of
      Left Dict -> Density n1 $ mat1 + mat2
      Right _   -> error $ "Cannot add mismatched matrices "
                            ++ show mat1 ++ " and " ++ show mat2
  _ - _ = undefined
  _ * _ = undefined
  abs _ = undefined
  signum = undefined
  fromInteger n = undefined

instance Show (DensityMonad a) where
  show = show . getDensity

-- Debugging help
m0 = square @Two [1,0,0,0
                 ,0,1,0,0
                 ,0,0,0,1
                 ,0,0,1,0]
rho = Density three $ \case
  (0,0) -> -0.5
  (0,2) ->  0.5
  (3,0) -> -0.5
  (3,2) ->  0.5
  (_,_) -> 0

{-
-- produces an association list
listToAssoc :: forall n. SingI n => [Int] -> [(BNat n, BNat n)]
listToAssoc ls = (fromIntegral <$> [0..n]) `zip` (fromIntegral <$> pad ls)
  where
    n = maxBNat (sing :: Sing n)

pad :: [Int] -> [Int]
pad ls = ls ++ ([0..] \\ ls)

mkPermutation :: [(BNat n,BNat n)] -> BNat n -> BNat n
mkPermutation ls i =
  case lookup i ls of
    Just i' -> i'
    Nothing -> error $ "Permutation " ++ show ls
                        ++ " not well defined at index " ++ show i

permute :: forall m. SingI m => [(BNat m, BNat m)] -> Matrix m m -> Matrix m m
permute ls mat (i,j) = mat (mkPermutation ls i, mkPermutation ls j)

expand :: forall m n. (SingI m, SingI n)
       => [(BNat (m `Times` n), BNat (m `Times` n))]
       -> Matrix m m -> Matrix (m `Times` n) (m `Times` n)
expand ls mat = Tr.trace ("Expanding..." ++ show mat) $
                withSingI (timesSNat (sing :: Sing m) (sing :: Sing n)) $
                  permute ls $ mat `kron` ident @n

applyMatrix :: forall m p. (SingI m, SingI p, (p < m) ~ 'False)
            => Squared m -> [Int] -> Squared p -> Squared p
applyMatrix mat ls ρ =
  Tr.trace ("Applying " ++ show ls) $
  case raiseToPlus two m minus of {Dict ->
  case plusMinus m p of {Dict ->
  withSingI (raiseToSNat two m) $
  withSingI (raiseToSNat two minus) $
  withSingI (raiseToSNat two p) $
    let mat1 = expand @(Two `RaiseTo` m) @(Two `RaiseTo` (p `Minus` m))
                       (swap <$> ls') (transpose mat)
        mat2 = expand @(Two `RaiseTo` m) @(Two `RaiseTo` (p `Minus` m)) ls' mat
        res  = Tr.trace "Middle of application" $
               mat1 `mult` ρ `mult` mat2
    in mat1 `seq` mat2 `seq` res `seq` Tr.trace "End of application" res
  }}
  where
    ls' :: [(BNat (Two `RaiseTo` p), BNat (Two `RaiseTo` p))]
    ls' = withSingI (raiseToSNat two p) $ listToAssoc ls
    m = (sing :: Sing m)
    p = (sing :: Sing p)
    minus :: Sing (p `Minus` m)
    minus = minusSNat p m

applyMatrixD :: forall m. SingI m => Squared m -> [Int] -> Density -> Density
applyMatrixD mat perm (Density (p :: Sing p) ρ) =
  Tr.trace ("Applying " ++ mat' ++ " to qubits " ++ show perm ++ " of \n" ++ show (Density p ρ)) $
  case compareSNat p (sing :: Sing m) of
    -- p < m
    Left Dict -> error "Cannot apply matrix to state: matrix too large"
    -- m <= p
    Right Dict ->
      let res = withSingI p $ Density p $ applyMatrix @m @p mat perm ρ
      in  res `seq` Tr.trace "Done!" res
  where
    mat' = withSingI (two `raiseToSNat` (sing :: Sing m)) $ show mat

applyUnitaryM :: forall m. SingI m => Squared m -> [Int] -> DensityMonad ()
applyUnitaryM mat ls = get >>= \ρ -> put $ applyMatrixD @m mat ls ρ

measM :: Int -> DensityMonad Bool
measM i = do
    ρ <- get
    (b,ρ') <- lift [ (False, applyMatrixD @One density0 [i] ρ)
                   , (True,  applyMatrixD @One density1 [i] ρ) ]
    put ρ'
    return b

-}


module StatReport (statReport, showStatEntryValue) where

import Fmt
import Data.Text (Text)
import Data.Fixed (showFixed)
import QuoteData (QField(..))
import Statistics (Statistic(..), StatEntry(..), StatQFieldData, StatInfo)

instance Buildable Statistic where
  build Mean = "Mean"
  build Min = "Minimum"
  build Max = "Maximum"
  build Days = "Days between Min/Max"

showStatEntryValue :: StatEntry -> String
showStatEntryValue StatEntry {..} = showFixed (removeTrailings stat qfield) value where
  removeTrailings Days _ = True
  removeTrailings Min Volume = True
  removeTrailings Max Volume = True
  removeTrailings _ _ = False

instance Buildable StatEntry where
  build se@StatEntry {..} = "" +|stat|+ ": " +|showStatEntryValue se|+ ""

instance Buildable StatQFieldData where
  build (qf, stats) = nameF ("Statistic for " +||qf||+ "") $ unlinesF stats

statReport :: StatInfo -> Text
statReport = fmt . unlinesF


module GalaxyStats
where

import qualified Data.Map as M
import Text.Printf

import Galaxy
import DataFunction
import Statistics
import Libaddutil.Named
import Utils

numStarSystemsInGalaxy :: Galaxy a -> Int
numStarSystemsInGalaxy = M.size . starsystems

numStarsInSystem :: StarSystem a -> Int
numStarsInSystem = M.size . stars

numStarsInGalaxy :: Galaxy a -> Int
numStarsInGalaxy g = sum $ map numStarsInSystem ((M.elems . starsystems) g)

starsPerStarSystemInGalaxy :: Galaxy a -> Float
starsPerStarSystemInGalaxy g = fromIntegral (numStarsInGalaxy g) / fromIntegral (numStarSystemsInGalaxy g)

bodiesPerStarsInGalaxy g = fromIntegral (numBodiesInGalaxy g) / fromIntegral (numStarsInGalaxy g)
planetsPerStarsInGalaxy g = fromIntegral (numPlanetsInGalaxy g) / fromIntegral (numStarsInGalaxy g)
bodiesPerStarsystemsInGalaxy g = fromIntegral (numBodiesInGalaxy g) / fromIntegral (numStarSystemsInGalaxy g)
planetsPerStarsystemsInGalaxy g = fromIntegral (numPlanetsInGalaxy g) / fromIntegral (numStarSystemsInGalaxy g)
moonsPerBodiesInGalaxy g = fromIntegral (numMoonsInGalaxy g) / fromIntegral (numBodiesInGalaxy g - numMoonsInGalaxy g)

numBodiesOrbitingStar :: Star a -> Int
numBodiesOrbitingStar s =
  let ps = (M.elems . planets) s
  in length ps + sum (map numSatellites ps)

numSatellites :: Planet a -> Int
numSatellites = M.size . satellites

numBodiesInStarSystem :: StarSystem a -> Int
numBodiesInStarSystem = sum . map numBodiesOrbitingStar . M.elems . stars

numBodiesInGalaxy :: Galaxy a -> Int
numBodiesInGalaxy = sum . map numBodiesInStarSystem . M.elems . starsystems

moonsInGalaxy :: Galaxy a -> [Planet a]
moonsInGalaxy = concatMap M.elems . map satellites . allBodies

numMoonsInGalaxy :: Galaxy a -> Int
numMoonsInGalaxy = length . moonsInGalaxy

numStarsByType t g = (filter (\s -> spectralType s == t)) (allStars g)

numBStarsInGalaxy = length . numStarsByType SpectralTypeB
numAStarsInGalaxy = length . numStarsByType SpectralTypeA
numFStarsInGalaxy = length . numStarsByType SpectralTypeF
numGStarsInGalaxy = length . numStarsByType SpectralTypeG
numKStarsInGalaxy = length . numStarsByType SpectralTypeK
numMStarsInGalaxy = length . numStarsByType SpectralTypeM

numBodiesByPlanetType a g = length $ (filter (\p -> planettype p == a)) (allBodies g)

numRockyBodiesByAtmosphere :: Atmosphere -> Galaxy a -> Int
numRockyBodiesByAtmosphere a g = length $ filter (\p -> planettype p == RockyPlanet a) (allBodies g)

isGasGiant p = case planettype p of
                 SmallGasGiant     -> True
                 MediumGasGiant    -> True
                 LargeGasGiant     -> True
                 VeryLargeGasGiant -> True
                 _                 -> False

numGasGiantsInGalaxy g = length $ (filter isGasGiant) (allBodies g)

numWaterWeatherSystemBodiesInGalaxy :: Galaxy a -> Int
numWaterWeatherSystemBodiesInGalaxy = numRockyBodiesByAtmosphere WaterWeatherSystem
numNoAtmosphereBodiesInGalaxy g = numBodiesByPlanetType NoAtmosphere g + numBodiesByPlanetType Planetoid g
numHabitableBodiesInGalaxy :: Galaxy a -> Int
numHabitableBodiesInGalaxy g = length $ filter sustainsLife (allBodies g)

showPerc :: Int -> Int -> String
showPerc a b = printf " (%.2f%%)" (100.0 * (fromIntegral a / fromIntegral b) :: Float)

showstars :: String -> Int -> Int -> String
showstars c st tot = c ++ show st ++ (showPerc st tot) ++ "\n"

planetMassesInGalaxy g = map planetMass (allBodies g)

minBodyMassInGalaxy = minimum . planetMassesInGalaxy

maxBodyMassInGalaxy = maximum . planetMassesInGalaxy

medBodyMassInGalaxy = median . planetMassesInGalaxy

avgBodyMassInGalaxy = average . planetMassesInGalaxy

planetTemperaturesInGalaxy g = map planetTemperature (allPlanets g)

minBodyTemperatureInGalaxy = kelvinToCelsius . minimum . planetTemperaturesInGalaxy

maxBodyTemperatureInGalaxy = kelvinToCelsius . maximum . planetTemperaturesInGalaxy

avgBodyTemperatureInGalaxy = kelvinToCelsius . averageInt . planetTemperaturesInGalaxy

medBodyTemperatureInGalaxy = kelvinToCelsius . median . planetTemperaturesInGalaxy

numPlanetsInGalaxy g = numBodiesInGalaxy g - numMoonsInGalaxy g

galaxyStats :: Galaxy a -> String
galaxyStats g =
  let starstotal = numStarsInGalaxy g
      bodiestotal = numBodiesInGalaxy g
      bstars = numBStarsInGalaxy g
      astars = numAStarsInGalaxy g
      fstars = numFStarsInGalaxy g
      gstars = numGStarsInGalaxy g
      kstars = numKStarsInGalaxy g
      mstars = numMStarsInGalaxy g
      gasbodies = numGasGiantsInGalaxy g
      waterbodies = numWaterWeatherSystemBodiesInGalaxy g
      noatmbodies = numNoAtmosphereBodiesInGalaxy g
      habitablebodies = numHabitableBodiesInGalaxy g
  in
     "Galaxy name: " ++ name g ++ "\n" ++
     "Number of star systems: " ++ (show . numStarSystemsInGalaxy) g ++ "\n" ++
     "Number of stars: " ++ (show . numStarsInGalaxy) g ++ "\n" ++
     "Number of stars / star system: " ++ (show2f . starsPerStarSystemInGalaxy) g ++ "\n" ++
     showstars "Number of stars of spectral type B: " bstars starstotal ++
     showstars "Number of stars of spectral type A: " astars starstotal ++
     showstars "Number of stars of spectral type F: " fstars starstotal ++
     showstars "Number of stars of spectral type G: " gstars starstotal ++
     showstars "Number of stars of spectral type K: " kstars starstotal ++
     showstars "Number of stars of spectral type M: " mstars starstotal ++
     "Number of bodies: " ++ (show . numBodiesInGalaxy) g ++ "\n" ++
     "Number of planets: " ++ (show . numPlanetsInGalaxy) g ++ "\n" ++
     "Number of planets / starsystem: " ++ (show2f . planetsPerStarsystemsInGalaxy) g ++ "\n" ++
     "Number of planets / star: " ++ (show2f . planetsPerStarsInGalaxy) g ++ "\n" ++
     "Number of moons: " ++ (show . numMoonsInGalaxy) g ++ "\n" ++
     "Number of moons / planet: " ++ (show2f . moonsPerBodiesInGalaxy) g ++ "\n" ++
     "Minimum body mass (Earths): " ++ (show . minBodyMassInGalaxy) g ++ "\n" ++
     "Median body mass (Earths): " ++ (show . medBodyMassInGalaxy) g ++ "\n" ++
     "Average body mass (Earths): " ++ (show . avgBodyMassInGalaxy) g ++ "\n" ++
     "Maximum body mass (Earths): " ++ (show . maxBodyMassInGalaxy) g ++ "\n" ++
     "Minimum body temperature (Celsius): " ++ (show . minBodyTemperatureInGalaxy) g ++ "\n" ++
     "Median body temperature (Celsius): " ++ (show . medBodyTemperatureInGalaxy) g ++ "\n" ++
     "Average body temperature (Celsius): " ++ (show . avgBodyTemperatureInGalaxy) g ++ "\n" ++
     "Maximum body temperature (Celsius): " ++ (show . maxBodyTemperatureInGalaxy) g ++ "\n" ++
     showstars "Number of gas giants: " gasbodies bodiestotal ++
     showstars "Number of oxygen bodies: " waterbodies bodiestotal ++
     showstars "Number of bodies without atmosphere: " noatmbodies bodiestotal ++
     showstars "Number of habitable bodies: " habitablebodies bodiestotal


module Main where

import Control.Monad

import Criterion.Main

import Options.Applicative

import Statistics.Quantile.Bench
import Statistics.Quantile.Exact
import Statistics.Quantile.Approximate.Sampling
import Statistics.Quantile.Bench.Gen
import Statistics.Quantile.Types
import Statistics.Quantile.Util

import System.Environment

import System.IO

data Params = Params Bool
  deriving (Eq, Show)

main :: IO ()
main = do
  hSetBuffering stdin LineBuffering
  hSetBuffering stdout LineBuffering
  (fp, fh) <- openTempFile tmpDir "bench"
  void $ countOut fh stream
  hClose fh
  s <- lookupEnv "BENCH_SIZE"
  defaultMain $
    [ benchSelector [benchSize s] "median/exact" median external fp
    , benchBufferingSelector sqrtN [benchSize s] "median/sampling-sqrt(n)" median sampling fp
    , benchBufferingSelector sqrtN [benchSize s] "median/jackknife-sampling-sqrt(n)" median samplingJackknife fp
    ]
  where stream = streamHandle stdin

        benchSize = maybe 10000 read
--        stream = Stream $ gauss 3.14159 10000


{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -Wcompat #-}
{-# OPTIONS_GHC -Wincomplete-record-updates #-}
{-# OPTIONS_GHC -Wincomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wredundant-constraints #-}

module CloudSunlight
  ( cloudSunlightExperiment,
  )
where

import Control.Monad
import qualified Data.ByteString.Lazy as BL
import Data.Csv
import qualified Data.List as L
import qualified Data.Vector.Unboxed as V
import qualified Statistics.Regression as R
import qualified Statistics.Sample as S
import System.Random.MWC
import qualified System.Random.MWC.Distributions as D

meanToa :: Double
meanToa = 340.0

toa :: GenIO -> IO Double
toa = D.normal meanToa 30.0

-- | magnitude of aerosol affect on sunlight, between 1.0 and 0.0 (1.0 = max effect)
aerosolAffectSunlight :: Double
aerosolAffectSunlight = 1.0

-- | factor that cloud reduces sunlight, if there is cloud
cloudAffectSunlight :: Double
cloudAffectSunlight = 0.6

sunlight ::
  -- | non-dim measure of aerosol
  Double ->
  -- | cloud
  Bool ->
  -- | U_sun (i.e. TOA)
  Double ->
  -- | W/m2
  Double
sunlight a c toa' = if c then clearSky * cloudAffectSunlight else clearSky
  where
    clearSky = toa' * (1.0 - aerosolAffectSunlight * a)

cloud :: Double -> Double -> Bool
cloud a u
  | (a + u) > 1.0 = True
  | otherwise = False

aerosol :: Double -> Double
aerosol u = u

genTuple :: GenIO -> IO (Double, String, Double)
genTuple gen = do
  uAero <- uniform gen
  uCld <- uniform gen
  uSun <- toa gen
  let a = aerosol uAero
  let c = cloud a uCld
  let s = sunlight a c uSun
  return (a, if c then "Cloudy" else "Clear", s)

getCloud :: (a, b, c) -> b
getCloud (_a, c, _s) = c

getSun :: (a, b, c) -> c
getSun (_a, _c, s) = s

getAerosol :: (a, b, c) -> a
getAerosol (a, _c, _s) = a

calcRegression :: V.Vector Double -> Double -> Double
calcRegression ols x = x * ols V.! 0 + ols V.! 1

calcStats ::
  [(Double, String, Double)] ->
  [(Double, String, Double)] ->
  (V.Vector Double, V.Vector Double, Double, Double)
calcStats cloud' clear =
  ( olsCloud,
    olsClear,
    avgRegress olsCloud - avgRegress olsClear,
    naiveCloud - naiveClear
  )
  where
    calcAvgs xs = (as, ols, S.mean ss)
      where
        ss = V.fromList (fmap getSun xs)
        as = V.fromList (fmap getAerosol xs)
        ols = fst . R.olsRegress [as] $ ss
    (aCloud, olsCloud, naiveCloud) = calcAvgs cloud'
    (aClear, olsClear, naiveClear) = calcAvgs clear
    totAs = aCloud <> aClear
    avgRegress ols = S.mean $ V.map (calcRegression ols) totAs

cloudSunlightExperiment :: Int -> IO (Double, Double, Double)
cloudSunlightExperiment n = do
  gen <- create
  xs <- replicateM n (genTuple gen)
  let (cloud', clear) = L.partition ((== "Cloudy") . getCloud) xs
  let (olsCloud, olsClear, estDiff, naiveDiff) = calcStats cloud' clear
  let regPoints = [0, 0.1 .. 1]
  f "dat/cloudy.dat" cloud'
  f "dat/clear.dat" clear
  write
    "aerosol\tsunlight\n"
    "dat/olsCloudy.dat"
    (zip regPoints (fmap (calcRegression olsCloud) regPoints))
  write
    "aerosol\tsunlight\n"
    "dat/olsClear.dat"
    (zip regPoints (fmap (calcRegression olsClear) regPoints))
  return (estDiff, naiveDiff, (0.5 * meanToa * (cloudAffectSunlight - 1.0)))
  where
    f = write "aerosol\tcloud\tsunlight\n"
    write header' fpath =
      BL.writeFile fpath . (header' <>)
        . encodeWith
          (defaultEncodeOptions {encDelimiter = 9})


{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}

module ScalarWave where

import Control.Exception (assert)
import Control.Monad.Loops
import qualified Data.Vector.Unboxed as U
import Data.VectorSpace
import Dual
import Equations
import Foreign
import qualified Numeric.LinearAlgebra as L
import SSSFN
import System.CPUTime

default (Int)

--------------------------------------------------------------------------------

bndw :: Floating a => Coord a -> a
bndw (Coord u v) = swamp * cos (pi / 4 * swk * t) * cos (pi / 4 * swk * r)
  where
    t = v + u
    r = v - u
    swk = 0.1
    swamp = 0.1

guess1 :: Floating a => Coord a -> State a
guess1 c = State 1 1 0

guess :: (L.Container L.Vector a, Floating a) => Grid (State a)
guess = gsample (\(u, v) -> guess1 (Coord u v))

--------------------------------------------------------------------------------

eqnInt1 :: Num a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnInt1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  case eqn of
    EqnG01 -> eqnG01 s r2s r1sr su sv suv
    EqnG22 -> eqnG22 s r2s r1sr su sv suv
    EqnSW2 -> eqnSW2 s r2s r1sr su sv suv
    EqnG00 -> eqnG00 c s r2s r1su su suu
    EqnG11 -> eqnG11 c s r2s r1sv sv svv

eqnOrig1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnOrig1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  case eqn of
    EqnG01 -> q s - 1
    EqnG22 -> g s - 1
    EqnSW2 -> w s - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

-- i == 0, u == -1
eqnBndU1 :: (Eq a, Floating a) => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnBndU1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  assert (u c == -1) $
    case eqn of
      EqnG01 -> eqnG11 c s r2s r1sv sv svv
      EqnG22 -> g s - 1
      EqnSW2 -> w s - bndw c
      EqnG00 -> 0
      EqnG11 -> 0

-- j == 0, v == -1
eqnBndV1 :: (Eq a, Floating a) => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnBndV1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  assert (v c == -1) $
    case eqn of
      EqnG01 -> eqnG00 c s r2s r1su su suu
      EqnG22 -> g s - 1
      EqnSW2 -> w s - bndw c
      EqnG00 -> 0
      EqnG11 -> 0

eqnEndU1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnEndU1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  case eqn of
    EqnG01 -> q s - 1
    EqnG22 -> g s - 1
    EqnSW2 -> w s - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

eqnEndV1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnEndV1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  case eqn of
    EqnG01 -> q s - 1
    EqnG22 -> g s - 1
    EqnSW2 -> w s - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

equation1 :: (Eq a, Floating a) => Coord Int -> Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
equation1 (Coord i j) =
  if
    | i == 0 && j == 0 -> eqnOrig1
    | i == 0 && j == np - 1 -> eqnEndU1
    | i == np - 1 && j == 0 -> eqnEndV1
    | i == 0 -> eqnBndU1
    | j == 0 -> eqnBndV1
    | True -> eqnInt1

--------------------------------------------------------------------------------

g2s :: Storable a => Grid a -> State (Grid a)
g2s (Grid xs) =
  let [g, q, w] = L.takesV (replicate 3 (np * np)) xs
   in State (Grid g) (Grid q) (Grid w)

s2g :: Storable a => State (Grid a) -> Grid a
s2g (State (Grid g) (Grid q) (Grid w)) = Grid (L.vjoin [g, q, w])

eqn1 ::
  forall a.
  (L.Field a, RealFloat a, U.Unbox a) =>
  Eqn ->
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqn1 eqn c s =
  Grid $
    L.fromList
      [ eqnRow i j
        | i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    r2s = (r2project #>) <$> s
    r1sr = (r1derivR #>) <$> s
    r1su = (r1derivU #>) <$> s
    r1sv = (r1derivV #>) <$> s
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: Int -> Int -> a
    eqnRow i j =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> r2s)
            (at <$> r1sr)
            (at <$> r1su)
            (at <$> r1sv)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

eqns ::
  forall a.
  (L.Field a, RealFloat a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqns c s =
  Grid $
    L.fromList
      [ eqnRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    r2s = (r2project #>) <$> s
    r1sr = (r1derivR #>) <$> s
    r1su = (r1derivU #>) <$> s
    r1sv = (r1derivV #>) <$> s
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: (Int, Int, Eqn) -> a
    eqnRow (i, j, eqn) =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> r2s)
            (at <$> r1sr)
            (at <$> r1su)
            (at <$> r1sv)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

eqns5 ::
  forall a.
  (L.Field a, RealFloat a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqns5 c s =
  Grid $
    L.fromList
      [ eqnRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2, EqnG00, EqnG11],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    r2s = (r2project #>) <$> s
    r1sr = (r1derivR #>) <$> s
    r1su = (r1derivU #>) <$> s
    r1sv = (r1derivV #>) <$> s
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: (Int, Int, Eqn) -> a
    eqnRow (i, j, eqn) =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> r2s)
            (at <$> r1sr)
            (at <$> r1su)
            (at <$> r1sv)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

op ::
  forall a.
  (L.Container L.Vector a, L.Field a, RealFloat a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Op a
op (Coord u v) s =
  Op $
    L.fromRows
      [ jacRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    r2s = (r2project #>) <$> s
    r1sr = (r1derivR #>) <$> s
    r1su = (r1derivU #>) <$> s
    r1sv = (r1derivV #>) <$> s
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    jacRow :: (Int, Int, Eqn) -> L.Vector a
    jacRow (i, j, eqn) =
      L.vjoin $
        [ let con :: Grid a -> Dual a
              der :: Grid a -> Dual a
              con x = Dual (x ! (i, j)) 0
              der x = Dual (x ! (i, j)) 1
              scon :: State (Grid a) -> State (Dual a)
              sder :: State (Grid a) -> State (Dual a)
              scon (State g q w) = State (con g) (con q) (con w)
              sder (State g q w) = case var of
                VarG -> State (der g) (con q) (con w)
                VarQ -> State (con g) (der q) (con w)
                VarW -> State (con g) (con q) (der w)
           in ( ( dual $
                    equation1
                      (Coord i j)
                      eqn
                      (Coord (con u) (con v))
                      (sder s)
                      (scon r2s)
                      (scon r1sr)
                      (scon r1su)
                      (scon r1sv)
                      (scon su)
                      (scon sv)
                      (scon suv)
                      (scon suu)
                      (scon svv)
                )
                  `L.scale` (getOp delta L.! (i * np + j))
              )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (sder r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r2project L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (sder r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivR L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (sder r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (sder r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (sder su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (sder sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (sder suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (sder suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (sder svv)
                          )
                            `L.scale` (getOp derivVV L.! (i * np + j))
                        )
          | var <- [VarG, VarQ, VarW]
        ]

op5 ::
  forall a.
  (L.Container L.Vector a, L.Field a, RealFloat a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Op a
op5 (Coord u v) s =
  Op $
    L.fromRows
      [ jacRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2, EqnG00, EqnG11],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    r2s = (r2project #>) <$> s
    r1sr = (r1derivR #>) <$> s
    r1su = (r1derivU #>) <$> s
    r1sv = (r1derivV #>) <$> s
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    jacRow :: (Int, Int, Eqn) -> L.Vector a
    jacRow (i, j, eqn) =
      L.vjoin $
        [ let con :: Grid a -> Dual a
              der :: Grid a -> Dual a
              con x = Dual (x ! (i, j)) 0
              der x = Dual (x ! (i, j)) 1
              scon :: State (Grid a) -> State (Dual a)
              sder :: State (Grid a) -> State (Dual a)
              scon (State g q w) = State (con g) (con q) (con w)
              sder (State g q w) = case var of
                VarG -> State (der g) (con q) (con w)
                VarQ -> State (con g) (der q) (con w)
                VarW -> State (con g) (con q) (der w)
           in ( ( dual $
                    equation1
                      (Coord i j)
                      eqn
                      (Coord (con u) (con v))
                      (sder s)
                      (scon r2s)
                      (scon r1sr)
                      (scon r1su)
                      (scon r1sv)
                      (scon su)
                      (scon sv)
                      (scon suv)
                      (scon suu)
                      (scon svv)
                )
                  `L.scale` (getOp delta L.! (i * np + j))
              )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (sder r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r2project L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (sder r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivR L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (sder r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (sder r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (sder su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (sder sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (sder suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (sder suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (sder svv)
                          )
                            `L.scale` (getOp derivVV L.! (i * np + j))
                        )
          | var <- [VarG, VarQ, VarW]
        ]

--------------------------------------------------------------------------------

getTime :: IO Double
getTime = do
  t <- getCPUTime
  return $ fromInteger t / 1.0e+12

main0 :: IO ()
main0 =
  do
    putStrLn $ show $ gmaxabs $ coordU @Double
    putStrLn $ show $ gmaxabs $ coordV @Double
    putStrLn $ show $ gmaxabs $ coordR @Double
    putStrLn $ show $ gmaxabs $ coordR1 @Double
    putStrLn $ show $ omaxabs $ zproject @Double
    putStrLn $ show $ omaxabs $ r2project @Double
    putStrLn $ show $ omaxabs $ r1derivR @Double
    putStrLn $ show $ omaxabs $ r1derivU @Double
    putStrLn $ show $ omaxabs $ r1derivV @Double
    putStrLn $ show $ omaxabs $ derivU @Double
    putStrLn $ show $ omaxabs $ derivV @Double
    putStrLn $ show $ omaxabs $ derivUV @Double
    putStrLn $ show $ omaxabs $ derivUU @Double
    putStrLn $ show $ omaxabs $ derivVV @Double

main :: IO ()
main =
  do
    putStrLn "Spherically Symmetric Scalar Field in double-Null coordinates"
    let coords = Coord coordU (coordV @Double)
    let ini = State (gmap g guess) (gmap q guess) (gmap w guess)
    vars0 <- return (s2g ini)
    let done (n, vars) =
          let res = eqns coords (g2s vars)
              maxabsres = gmaxabs res
           in isNaN maxabsres
                || isInfinite maxabsres
                || maxabsres < 1.0e-12
                || n >= 5
    let iter (n, vars) = do
          let res = eqns coords (g2s vars)
          let res5 = eqns5 coords (g2s vars)
          t <- getTime
          putStrLn $
            ("iter " ++ show n)
              ++ ("   time " ++ show (round t) ++ " sec")
              ++ ("   |res| " ++ show (gmaxabs res))
              ++ ("   |res5| " ++ show (gmaxabs res5))
          let jac = op coords (g2s vars)
          let dvars = solve jac res
          return (n + 1, vars ^-^ dvars)
    let done5 (n, vars) =
          let res5 = eqns5 coords (g2s vars)
              maxabsres5 = gmaxabs res5
           in isNaN maxabsres5
                || isInfinite maxabsres5
                || maxabsres5 < 1.0e-12
                || n >= 5
    let iter5 (n, vars) = do
          let res = eqns coords (g2s vars)
          let res5 = eqns5 coords (g2s vars)
          t <- getTime
          putStrLn $
            ("iter " ++ show n)
              ++ ("   time " ++ show (round t) ++ " sec")
              ++ ("   |res| " ++ show (gmaxabs res))
              ++ ("   |res5| " ++ show (gmaxabs res5))
          let jac5 = op5 coords (g2s vars)
          let dvars = solveLS jac5 res5
          return (n + 1, vars ^-^ dvars)
    -- (n, vars) <- iterateUntilM done iter (0, vars0)
    (n, vars) <- iterateUntilM done5 iter5 (0, vars0)
    let res = eqns coords (g2s vars)
    let res5 = eqns5 coords (g2s vars)
    t <- getTime
    putStrLn $
      ("iter " ++ show n)
        ++ ("   time " ++ show (round t) ++ " sec")
        ++ ("   |res| " ++ show (gmaxabs res))
        ++ ("   |res5| " ++ show (gmaxabs res5))
    putStrLn $ "var.g " ++ show (gmap approx $ g $ g2s vars)
    putStrLn $ "var.q " ++ show (gmap approx $ q $ g2s vars)
    putStrLn $ "var.w " ++ show (gmap approx $ w $ g2s vars)
    putStrLn $ "res.G01 " ++ show (gmap approx $ eqn1 EqnG01 coords (g2s vars))
    putStrLn $ "res.G22 " ++ show (gmap approx $ eqn1 EqnG22 coords (g2s vars))
    putStrLn $ "res.SW2 " ++ show (gmap approx $ eqn1 EqnSW2 coords (g2s vars))
    putStrLn $ "res.G00 " ++ show (gmap approx $ eqn1 EqnG00 coords (g2s vars))
    putStrLn $ "res.G11 " ++ show (gmap approx $ eqn1 EqnG11 coords (g2s vars))
    putStrLn $ "S var.g " ++ show (gmap approx $ inv vandermonde #> (g $ g2s vars))
    putStrLn $ "S var.q " ++ show (gmap approx $ inv vandermonde #> (q $ g2s vars))
    putStrLn $ "S var.w " ++ show (gmap approx $ inv vandermonde #> (w $ g2s vars))
    putStrLn $ "S res.G01 " ++ show (gmap approx $ inv vandermonde #> (eqn1 EqnG01 coords (g2s vars)))
    putStrLn $ "S res.G22 " ++ show (gmap approx $ inv vandermonde #> (eqn1 EqnG22 coords (g2s vars)))
    putStrLn $ "S res.SW2 " ++ show (gmap approx $ inv vandermonde #> (eqn1 EqnSW2 coords (g2s vars)))
    putStrLn $ "S res.G00 " ++ show (gmap approx $ inv vandermonde #> (eqn1 EqnG00 coords (g2s vars)))
    putStrLn $ "S res.G11 " ++ show (gmap approx $ inv vandermonde #> (eqn1 EqnG11 coords (g2s vars)))
    putStrLn "Done."
  where
    approx :: RealFrac a => a -> a
    approx x = fromInteger (round (1.0e+10 * x)) / 1.0e+10


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
module Main where

import Data.Complex
import Data.Monoid
import Control.Arrow ((***))
import Control.Monad.State
import Control.Monad.Identity
import System.Random
import Lucid.Svg

import Data.List (foldl')

-- |Defines the finite trace of a coalgebra up to n iterations
traceN :: (Monad m) => Int -> (x -> m [x]) -> x -> m [x]
traceN 0 c seed = return []
traceN 1 c seed = c seed
traceN n c seed = concat <$> (traceN (n-1) c seed >>= mapM c)

data NDFPP = F | T Bool | Push | Pop | X
  deriving (Eq , Show)

-- |Stochastic seaweed L-system
--
-- > F=FF
-- > X=F-[[X]+X]+F[+FX]-X
--    ,F+[[X]-X]-F[-FX]+X
--
stochweed :: (RandomGen g) => NDFPP -> State g [NDFPP]
stochweed F = return [F , F]
stochweed X = run <$> state random
  where
    run b = [F , T (not b) , Push , Push , X , Pop , T b , X , Pop
            ,T b , F , Push , T b , F , X , Pop , T (not b) , X]
stochweed x = return [x]

type Seed = Int

stochweedDet :: Seed -> Int -> [NDFPP]
stochweedDet seed n = evalState (traceN n stochweed X) (mkStdGen seed)

rot :: Float -> Complex Float -> Complex Float
rot angle = (* cis angle)

{-
render :: Float -> [NDFPP] -> State [Complex Float] (Svg ())
render angle = foldl' (\r x -> r >>= go x) (return $ return ())
  where
    go :: NDFPP -> Svg () -> State [Complex Float] (Svg ())
    go F = path_ [ d_ (mA
-}

mvTo   c = mA (realPart c) (imagPart c)
lineTo c = lA (realPart c) (imagPart c)

-- state is current dir and current pos
renderS :: Float -> [NDFPP] -> StateT [(Complex Float , Complex Float)] Svg ()
renderS angle []     = return ()
renderS angle (i:is) = render1 i >> renderS angle is
  where
    render1 :: NDFPP -> StateT [(Complex Float , Complex Float)] Svg ()
    render1 i = do
      (dir , p):ps <- get
      case i of
        T True  -> put ((rot   angle  dir , p):ps)
        T False -> put ((rot (-angle) dir , p):ps)
        Push -> put ((dir , p):(dir , p):ps)
        Pop  -> put ps
        _ -> let p' = dir + p
              in put ((dir, p'):ps)
              >> lift (path_ [ d_ (mvTo p <> lineTo p' <> z) , stroke_ "black" ])

svg :: Svg () -> Svg ()
svg content = do
  doctype_
  with (svg11_ content) [version_ "1.1", width_ "300" , height_ "200"]

main :: IO ()
main = let tree = stochweedDet 123513 3
        in print $ svg $ evalStateT (renderS 22.5 tree) [(0 :+ 10 , mkPolar 0 0)]


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
#if __GLASGOW_HASKELL__ >= 800
  {-# OPTIONS_GHC -Wno-redundant-constraints #-}
#endif
-- |
-- Module      : Graphics.Image.ColorSpace.Complex
-- Copyright   : (c) Alexey Kuleshevich 2016
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Graphics.Image.ColorSpace.Complex (
  -- *** Rectangular form
  Complex(..), (+:), realPart, imagPart,
  -- *** Polar form
  mkPolar, cis, polar, magnitude, phase,
  -- *** Conjugate
  conjugate
  ) where

import Graphics.Image.Interface (Pixel)
import Control.Applicative
import Data.Complex (Complex(..))
import qualified Data.Complex as C hiding (Complex(..))

infix 6 +:

-- | Constrcut a complex pixel from two pixels representing real and imaginary parts.
--
-- @ PixelRGB 4 8 6 '+:' PixelRGB 7 1 1 __==__ PixelRGB (4 ':+' 7) (8 ':+' 1) (6 ':+' 1) @
--
(+:) :: Applicative (Pixel cs) => Pixel cs e -> Pixel cs e -> Pixel cs (Complex e)
(+:) = liftA2 (:+)
{-# INLINE (+:) #-}

-- | Extracts the real part of a complex pixel.
realPart :: (Applicative (Pixel cs), RealFloat e) => Pixel cs (Complex e) -> Pixel cs e
realPart = liftA C.realPart
{-# INLINE realPart #-}

-- | Extracts the imaginary part of a complex pixel.
imagPart :: (Applicative (Pixel cs), RealFloat e) => Pixel cs (Complex e) -> Pixel cs e
imagPart = liftA C.imagPart
{-# INLINE imagPart #-}

-- | Form a complex pixel from polar components of magnitude and phase.
mkPolar :: (Applicative (Pixel cs), RealFloat e) =>
           Pixel cs e -> Pixel cs e -> Pixel cs (Complex e)
mkPolar = liftA2 C.mkPolar
{-# INLINE mkPolar #-}

-- | @'cis' t@ is a complex pixel with magnitude 1 and phase t (modulo @2*'pi'@).
cis :: (Applicative (Pixel cs), RealFloat e) => Pixel cs e -> Pixel cs (Complex e)
cis = liftA C.cis
{-# INLINE cis #-}

-- | The function @'polar'@ takes a complex pixel and returns a (magnitude, phase)
-- pair of pixels in canonical form: the magnitude is nonnegative, and the phase
-- in the range @(-'pi', 'pi']@; if the magnitude is zero, then so is the phase.
polar :: (Applicative (Pixel cs), RealFloat e) => Pixel cs (Complex e) -> (Pixel cs e, Pixel cs e)
polar !zPx = (magnitude zPx, phase zPx)
{-# INLINE polar #-}

-- | The nonnegative magnitude of a complex pixel.
magnitude :: (Applicative (Pixel cs), RealFloat e) => Pixel cs (Complex e) -> Pixel cs e
magnitude = liftA C.magnitude
{-# INLINE magnitude #-}

-- | The phase of a complex pixel, in the range @(-'pi', 'pi']@. If the
-- magnitude is zero, then so is the phase.
phase :: (Applicative (Pixel cs), RealFloat e) => Pixel cs (Complex e) -> Pixel cs e
phase = liftA C.phase
{-# INLINE phase #-}

-- | The conjugate of a complex pixel.
conjugate :: (Applicative (Pixel cs), RealFloat e) => Pixel cs (Complex e) -> Pixel cs (Complex e)
conjugate = liftA C.conjugate
{-# INLINE conjugate #-}