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

module Main
  ( main

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

-- base
import Data.Coerce
  ( coerce )
import Data.Foldable
  ( for_ )
import GHC.Exts
  ( Proxy#, proxy# )
import GHC.Generics
  ( Generic )
import GHC.TypeNats
  ( type (-) )

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

-- brush-strokes
import Calligraphy.Brushes
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
      ]

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])


-}

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

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 )