{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}

module Main
  ( main

    -- Testing
  , TestCase(..)
  , testCases
  , testCaseStrokeFunctions
  , eval
  , mkVal, mkBox
  , potentialCusp
  , dEdsdcdt
  )
  where

-- base
import Control.Concurrent.MVar
  ( newMVar )
import Data.Coerce
  ( coerce )
import Data.Foldable
  ( for_ )
import Data.List
  ( intercalate )
import GHC.Exts
  ( Proxy#, proxy# )
import GHC.Generics
  ( Generic )
import GHC.TypeNats
  ( type (-) )
import Numeric
  ( showFFloat )

-- containers
import Data.Sequence
  ( Seq )
import qualified Data.Sequence as Seq
  ( index )
import Data.Tree
  ( foldTree )

-- brush-strokes
import Calligraphy.Brushes
import Debug.Utils
  ( logToFile )
import Math.Algebra.Dual
import Math.Bezier.Spline
import Math.Bezier.Stroke
import Math.Bezier.Stroke.EnvelopeEquation
import Math.Differentiable
import Math.Interval
import Math.Linear
import Math.Module
import Math.Ring
  ( Transcendental )

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

main :: IO ()
main = for_ testCases $ \ testCase@( TestCase { testName, testAlgorithmParams } ) -> do
  let ( _, testStrokeFnI ) = testCaseStrokeFunctions testCase
      ( newtTrees, ( dunno, sols ) ) = computeCusps testAlgorithmParams testStrokeFnI
      showedTrees = map ( uncurry showIntervalNewtonTree ) newtTrees
  putStrLn $ unlines $
    [ ""
    , "Test case '" ++ testName ++ "':" ] ++
    map ( "    " ++ )
      [ " #sols: " ++ show (length sols)
      , "#dunno: " ++ show (length dunno)
      , "#trees: " ++ show @Int (sum @_ @Int $ map (foldTree ( \ _ bs -> 1 + sum bs )) showedTrees)
      , " dunno: " ++ show dunno
      , "  sols: " ++ show sols
      ]
  --logFileMVar <- newMVar "logs/trickyCusp.log"
  --logToFile logFileMVar (unlines logLines)
  --  `seq` return ()

testCases :: [ TestCase ]
testCases = [ ellipse , trickyCusp2 ]

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

data TestCase =
  forall nbParams. ParamsCt nbParams =>
  TestCase
  { testName   :: !String
  , testBrush  :: !( Brush nbParams )
  , testStroke :: !( Point nbParams, Curve Open () ( Point nbParams ))
  , testAlgorithmParams :: !CuspAlgorithmParams
  }

testCaseStrokeFunctions
  :: TestCase
  -> ( ℝ 1 -> Seq ( ℝ 1 -> StrokeDatum 2 () )
     , 𝕀ℝ 1 -> Seq ( 𝕀ℝ 1 -> StrokeDatum 3 𝕀 ) )
testCaseStrokeFunctions ( TestCase { testStroke = ( sp0, crv ), testBrush } ) =
  getStrokeFunctions testBrush sp0 crv

-- Utilities to use in GHCi to help debugging.

eval
  :: ( I i ( ℝ 1 ) -> Seq ( I i ( ℝ 1 ) -> StrokeDatum k i ) )
  -> ( I i ( ℝ 1 ), Int, I i ( ℝ 1 ) )
  -> StrokeDatum k i
eval f ( t, i, s ) = ( f t `Seq.index` i ) s

mkVal :: Double -> Int -> Double -> ( ℝ 1, Int, ℝ 1 )
mkVal t i s = ( ℝ1 t, i, ℝ1 s )

mkBox :: ( Double, Double ) -> Int -> ( Double, Double ) -> Box
mkBox ( t_min, t_max ) i ( s_min, s_max ) =
  ( 𝕀 ( ℝ1 t_min ) ( ℝ1 t_max ) , i, 𝕀 ( ℝ1 s_min ) ( ℝ1 s_max ) )

potentialCusp :: StrokeDatum 3 𝕀 -> Bool
potentialCusp
  ( StrokeDatum
    { ee = D22 { _D22_v = 𝕀 ( ℝ1 ee_min ) ( ℝ1 ee_max ) }
    , 𝛿E𝛿sdcdt = D12 { _D12_v = T ( 𝕀 ( ℝ2 vx_min vy_min ) ( ℝ2 vx_max vy_max ) )}
    }
  ) = ee_min <= 0 && ee_max >= 0
    && vx_min <= 0 && vx_max >= 0
    && vy_min <= 0 && vy_max >= 0

dEdsdcdt :: StrokeDatum k i -> D ( k - 2 ) ( I i ( ℝ 2 ) ) ( T ( I i ( ℝ 2 ) ) )
dEdsdcdt ( StrokeDatum { 𝛿E𝛿sdcdt = v } ) = v

{-
let (f, fI) = testCaseStrokeFunctions trickyCusp2

take 10 $ Data.List.sortOn ( \ ( _, ℝ1 e, v) -> abs e + norm v ) [ let { v = mkVal x 3 y; d = eval f v } in ( v, _D12_v $ ee d, _D0_v $ dEdsdcdt d ) | x <- [0.57,0.5701 .. 0.58], y <- [0.29,0.291..0.3] ]
> ((ℝ1 0.5798800000000057,3,ℝ1 0.267980000000008),ℝ1 -2.8596965543670194e-4,V2 7.79559474412963e-2 2.0389671921293484e-2)

potentialCusp $ eval fI $ mkBox (0.5798, 0.5799) 3 (0.26798, 0.26799)
> True

let nbPotentialSols b = let ( _newtTrees, ( dunno, sols ) ) = intervalNewtonGSFrom NoPreconditioning 1e-7 fI b in length dunno + length sols

nbPotentialSols $ mkBox (0.5798, 0.5799) 3 (0.26798, 0.26799)
1

nbPotentialSols $ mkBox (0.5798, 0.675) 3 (0.26798, 0.26799)
0

let showTrees b = map ( uncurry showIntervalNewtonTree ) $ fst $ intervalNewtonGSFrom NoPreconditioning 1e-7 fI b

putStrLn $ unlines $ map Data.Tree.View.showTree $ showTrees $ mkBox (0.5798, 0.675) 3 (0.26798, 0.26799)

([ℝ1 0.5798, ℝ1 0.675],3,[ℝ1 0.26798, ℝ1 0.26799]) (area 0.000001) N []
 └─ ([ℝ1 0.5973000285624527, ℝ1 0.6750000000000002],3,[ℝ1 0.26798, ℝ1 0.26799000000000006]) (area 0.000001) NoSolution "ee" ([ℝ1 0.5973000285624527, ℝ1 0.6750000000000002],3,[ℝ1 0.26798, ℝ1 0.26799000000000006])

eval fI $ mkBox (0.5798, 0.675) 3 (0.26798, 0.26799)
> D12 {_D12_v = T[ℝ2 -10088.6674944889 -3281.3820867312834, ℝ2 4124.668381545453 4524.807156085763], _D12_dx = TT[ℝ2 -173746.97965005718 -33281.18494907289, ℝ2 298.2609121556852 23639.772884799597], _D12_dy = TT[ℝ2 -18454.27716258352 -28337.509817580823, ℝ2 1163.6949532017436 -13936.383137525536]}}
i.e.
> f = [ℝ2 -10088.6674944889 -3281.3820867312834, ℝ2 4124.668381545453 4524.807156085763]
> f_t = [ℝ2 -173746.97965005718 -33281.18494907289, ℝ2 298.2609121556852 23639.772884799597]
> f_s = [ℝ2 -18454.27716258352 -28337.509817580823, ℝ2 1163.6949532017436 -13936.383137525536]

(f, fI) = testCaseStrokeFunctions trickyCusp2
t = 𝕀 (ℝ1 0.5798) (ℝ1 0.675)
s = 𝕀 (ℝ1 0.26798) (ℝ1 0.26799)
t_mid = 0.5 * ( 0.5798 + 0.675 )
s_mid = 0.5 * ( 0.26798 + 0.26799 )
D12 ( T f ) ( T ( T f_t ) ) ( T ( T f_s ) ) = dEdsdcdt $ eval fI (t, 3, s)
t' = coerce ( (-) @( 𝕀 Double ) ) t ( singleton ( ℝ1 t_mid ) ) :: 𝕀ℝ 1
s' = coerce ( (-) @( 𝕀 Double ) ) s ( singleton ( ℝ1 s_mid ) ) :: 𝕀ℝ 1
a = ( f_t, f_s )
b = negV2 $ singleton $ midV2 f
[((t2', s2'), isContr)] = gaussSeidel a b (t', s')
t2 = coerce ( (+) @( 𝕀 Double ) ) t2' ( singleton ( ℝ1 t_mid ) ) :: 𝕀ℝ 1
s2 = coerce ( (+) @( 𝕀 Double ) ) s2' ( singleton ( ℝ1 s_mid ) ) :: 𝕀ℝ 1

t2
> [ℝ1 0.6102365832093095, ℝ1 0.6750000000000002]
s2
> [ℝ1 0.26798, ℝ1 0.26799000000000006]


let ( 𝕀 ( ℝ2 a11_lo a21_lo ) ( ℝ2 a11_hi a21_hi ), 𝕀 ( ℝ2 a12_lo a22_lo ) ( ℝ2 a12_hi a22_hi ) ) = a
let ( 𝕀 ( ℝ2 b1_lo b2_lo ) ( ℝ2 b1_hi b2_hi ) ) = b
let ( 𝕀 ( ℝ1 x1_lo ) ( ℝ1 x1_hi ), 𝕀 ( ℝ1 x2_lo ) ( ℝ1 x2_hi ) ) = ( t', s' )

a11 = 𝕀 a11_lo a11_hi
a12 = 𝕀 a12_lo a12_hi
a21 = 𝕀 a21_lo a21_hi
a22 = 𝕀 a22_lo a22_hi
b1  = 𝕀  b1_lo  b1_hi
b2  = 𝕀  b2_lo  b2_hi
x1  = 𝕀  x1_lo  x1_hi
x2  = 𝕀  x2_lo  x2_hi

( b1 - a12 * x2 )
> [2981.90728508591, 2982.0918278575364]

extendedRecip a11

-}

negV2 :: 𝕀ℝ 2 -> 𝕀ℝ 2
negV2 ( 𝕀 ( ℝ2 x_lo y_lo ) ( ℝ2 x_hi y_hi ) ) =
  let !( 𝕀 x'_lo x'_hi ) = negate $ 𝕀 x_lo x_hi
      !( 𝕀 y'_lo y'_hi ) = negate $ 𝕀 y_lo y_hi
  in 𝕀 ( ℝ2 x'_lo y'_lo ) ( ℝ2 x'_hi y'_hi )

midV2 :: 𝕀ℝ 2 -> ℝ 2
midV2 ( 𝕀 ( ℝ2 x_lo y_lo ) ( ℝ2 x_hi y_hi ) ) =
  ℝ2 ( 0.5 * ( x_lo + x_hi ) ) ( 0.5 * ( y_lo + y_hi ) )

logLines :: [ String ]
logLines =
  [ "f = dE/ds * dc/dt: f, df/dt, df/ds"
  , "{" ++
    (intercalate ","
     [ "{" ++ showD (midPoint t) ++ "," ++ showD (midPoint s) ++ ",{" ++ intercalate "," vals ++ "}}"
     | t <- map ( around 0.5798  ) [-0.05, -0.049.. 0.05]
     , let i = 3
     , s <- map ( around 0.26798 ) [-0.05, -0.049.. 0.05]
     , let StrokeDatum
             { 𝛿E𝛿sdcdt = D12 (T f) (T (T f_t)) (T (T f_s))
             } = (curvesI t `Seq.index` i) s
           ℝ2 vx vy = midPoint2 f
           ℝ2 vx_t vy_t = midPoint2 f_t
           ℝ2 vx_s vy_s = midPoint2 f_s
           vals = [ "{" ++ showD vx ++ "," ++ showD vy ++ "}"
                  , "{" ++ showD vx_t ++ "," ++ showD vy_t ++ "}"
                  , "{" ++ showD vx_s ++ "," ++ showD vy_s ++ "}"
                  ]
     ]
    ) ++ "}"
  ]
  where
    around :: Double -> Double -> 𝕀ℝ 1
    around z0 z = 𝕀 ( ℝ1 ( z + z0 - 1e-6 ) ) ( ℝ1 ( z + z0 + 1e-6 ) )
    ( _, curvesI ) = testCaseStrokeFunctions trickyCusp2
    midPoint (𝕀 (ℝ1 lo) (ℝ1 hi)) = 0.5 * ( lo + hi )
    midPoint2 (𝕀 (ℝ2 lo_x lo_y) (ℝ2 hi_x hi_y))
      = ℝ2 ( 0.5 * ( lo_x + hi_x ) ) ( 0.5 * ( lo_y + hi_y ) )

showD :: Double -> String
showD float = showFFloat (Just 6) float ""

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

ellipse :: TestCase
ellipse =
  TestCase
    { testName = "ellipse"
    , testBrush = ellipseBrush
    , testStroke = ( p0, LineTo ( NextPoint p1 ) () )
    , testAlgorithmParams =
        CuspAlgorithmParams
         { preconditioning = NoPreconditioning
         , maxWidth = 1e-7
         }
    }
  where
    mkPt x y w h phi =
      Point
        { pointCoords = ℝ2 x y
        , pointParams = Params $ ℝ3 w h phi
        }
    p0 = mkPt 0 0 10 25 0
    p1 = mkPt 100 0 15 40 pi

trickyCusp2 :: TestCase
trickyCusp2 =
  TestCase
    { testName = "trickyCusp2"
    , testBrush = circleBrush
    , testStroke = ( p0, Bezier3To p1 p2 ( NextPoint p3 ) () )
    , testAlgorithmParams =
        CuspAlgorithmParams
         { preconditioning = NoPreconditioning
         , maxWidth = 1e-7
         }
    }
  where
    mkPt x y =
      Point
        { pointCoords = ℝ2 x y
        , pointParams = Params $ ℝ1 5.0
        }
    p0 = mkPt 5e+1 -5e+1
    p1 = mkPt -7.72994362904069e+1 -3.124468786098509e+1
    p2 = mkPt -5.1505430313958364e+1 -3.9826386521527986e+1
    p3 = mkPt -5e+1 -5e+1

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

type ParamsCt nbParams
  = ( Show ( ℝ nbParams )
    , HasChainRule Double 2 ( ℝ nbParams )
    , HasChainRule ( 𝕀 Double ) 3 ( 𝕀 ( ℝ nbParams ) )
    , Applicative ( D 2 ( ℝ nbParams ) )
    , Applicative ( D 3 ( ℝ nbParams ) )
    , Traversable ( D 2 ( ℝ nbParams ) )
    , Traversable ( D 3 ( ℝ nbParams ) )
    , Representable Double ( ℝ nbParams )
    , Module Double ( T ( ℝ nbParams ) )
    , Module ( 𝕀 Double ) ( T ( 𝕀 ( ℝ nbParams ) ) )
    , Module ( D 2 ( ℝ nbParams ) Double ) ( D 2 ( ℝ nbParams ) ( ℝ 2 ) )
    , Module ( D 3 ( ℝ nbParams ) ( 𝕀 Double ) ) ( D 3 ( ℝ nbParams ) ( 𝕀 ( ℝ 2 ) ) )
    , Transcendental ( D 2 ( ℝ nbParams ) Double )
    , Transcendental ( D 3 ( ℝ nbParams ) ( 𝕀 Double ) )
    )

newtype Params nbParams = Params { getParams :: ( ℝ nbParams ) }
deriving newtype instance Show ( ℝ nbParams ) => Show ( Params nbParams )

data Point nbParams =
  Point
    { pointCoords :: !( ℝ 2 )
    , pointParams :: !( Params nbParams ) }
  deriving stock Generic
deriving stock instance Show ( ℝ nbParams ) => Show ( Point nbParams )

data CuspAlgorithmParams =
  CuspAlgorithmParams
    { preconditioning :: !Preconditioner
    , maxWidth :: !Double
    }
  deriving stock Show

type Brush nbParams
  = forall {t} k (i :: t)
  .  DiffInterp k i ( ℝ nbParams )
  => Proxy# i
  -> ( forall a. a -> I i a )
  -> C k ( I i ( ℝ nbParams ) )
         ( Spline Closed () ( I i ( ℝ 2 ) ) )

getStrokeFunctions
  :: forall nbParams
  .  ParamsCt nbParams
  => Brush nbParams
      -- ^ brush shape
  -> Point nbParams
      -- ^ start point
  -> Curve Open () ( Point nbParams )
      -- ^ curve points
  -> ( ℝ 1 -> Seq ( ℝ 1 -> StrokeDatum 2 () )
     , 𝕀ℝ 1 -> Seq ( 𝕀ℝ 1 -> StrokeDatum 3 𝕀 ) )
getStrokeFunctions brush sp0 crv =
  let
    usedParams :: C 2 ( ℝ 1 ) ( ℝ nbParams )
    path :: C 2 ( ℝ 1 ) ( ℝ 2 )
    ( path, usedParams ) =
      pathAndUsedParams @2 @() coerce id ( getParams . pointParams )
        sp0 crv
    usedParamsI :: C 3 ( 𝕀ℝ 1 ) ( 𝕀ℝ nbParams )
    pathI :: C 3 ( 𝕀ℝ 1 ) ( 𝕀ℝ 2 )
    ( pathI, usedParamsI ) =
      pathAndUsedParams @3 @𝕀 coerce singleton ( getParams . pointParams )
        sp0 crv
  in ( brushStrokeData @2 @( ℝ nbParams ) coerce coerce
        path usedParams $
        brush @2 @() proxy# id
     , brushStrokeData @3 @( ℝ nbParams ) coerce coerce
        pathI usedParamsI $
        brush @3 @𝕀 proxy# singleton )
{-# INLINEABLE getStrokeFunctions #-}

computeCusps
  :: CuspAlgorithmParams
  -> ( 𝕀ℝ 1 -> Seq ( 𝕀ℝ 1 -> StrokeDatum 3 𝕀 ) )
  -> ( [ ( Box, IntervalNewtonTree Box ) ], ( [ Box ], [ Box ] ) )
computeCusps params =
  intervalNewtonGS ( preconditioning params ) ( maxWidth params )