{-# LANGUAGE StandaloneDeriving #-}

-- | Temporary hackery until all base libraries understand newtype Energy. And
-- yes, for testing too.

module BioInf.RNAFold.EnergyInt
  ( FoldFunctions (..)
  , Fold (..)
  ) where

import Biobase.Types.Ring
import Biobase.RNA
import Biobase.Constants
import Biobase.Structure.DotBracket
import Biobase.Structure

import Data.PrimitiveArray

import BioInf.RNAFold
import BioInf.RNAFold.Functions

-- for testing only
{-
import Biobase.Vienna.Default
import Debug.Trace.Tools
import Data.List.Split
import Text.Printf
import Control.Monad
import Biobase.Turner.Tables
-}


instance Ring Int where
  (.+.) = min
  (.*.) = (+)
  neg = negate
  one = 0
  zero = eInf
  isZero x = x>eMax
  n .^. k = n * k
  (.^^.) = error "write me"
  {-# INLINE (.+.) #-}
  {-# INLINE (.*.) #-}
  {-# INLINE neg #-}
  {-# INLINE one #-}
  {-# INLINE zero #-}
  {-# INLINE isZero #-}
  {-# INLINE (.^.) #-}

instance FoldFunctions Int where
  calcTermAU tAU p = if p/=vpCG && p/=vpGC then tAU else 0
  calcNinio maxNno nno l = min maxNno (nno * l)
  calcLargeLoop l = floor $ 108.856 * log (fromIntegral l / 30)
  {-# INLINE calcTermAU #-}
  {-# INLINE calcNinio #-}
  {-# INLINE calcLargeLoop #-}

-- TODO detach backtrack so that all instances that use this kind of
-- backtracking can immediately use it!

instance Fold Int where
  backtrack trnr inp (weak,strong,mbr1,mbr,extern) delta = filter ((<=0) . snd) $ map f $ ext 0 n delta where
    f (xs,z) = (Secondary (n+1) xs,optE+delta-z)
    n = snd $ bounds inp
    optE = extern!(0,n)
    newD d here next = d - (next - here)
    --
    -- All exterior loop decompositions
    --
    ext i j d = let ehere = extern!(i,j) in
      [ (x,z) -- shorter ext
      | i<j
      , let d' = newD d ehere (extern!(i+1,j))
      , d'>=0
      , (x,z) <- ext (i+1) j d'
      ] ++
      [ (x,z) -- loop at (i,j)
      | i<j
      , (_,enext) <- externalLoopIdx trnr inp strong i j -- _ = (i,j)
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z) <- str i j d'
      ] ++
      [ (x++y,z)
      | i<j
      , (k,enext) <- externalAddLoopIdx trnr inp strong extern i j
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z') <- str i k d'
      , (y,z) <- ext (k+1) j z' ++ [([],z') | z'>=0 && extern!(k+1,j)==0]
      ]
    --
    -- A strong stem on top of strong of weak
    --
    str i j d = let ehere = strong!(i,j) in
      [ ((i,j):x,z)
      | i<j
      , (_,enext) <- stackIdx trnr inp strong i j
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z) <- str (i+1) (j-1) d'
      ] ++
      [ ((i,j):x,z)
      | i<j
      , (_,enext) <- stackIdx trnr inp weak i j
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z) <- wea (i+1) (j-1) d'
      ]
    wea i j d = let ehere = weak!(i,j) in
      --
      -- A hairpin closes at (i,j)
      --
      [ ([(i,j)],d')
      | i<j
      , enext <- hairpinIdx trnr inp i j
      , let d' = newD d ehere enext
      , d'>=0
      ] ++
      --
      -- interior loops
      --
      [ ((i,j):x,z)
      | i<j
      , ((k,l),enext) <- (
          largeInteriorLoopIdx trnr inp strong i j ++
          tabbedInteriorLoopIdx trnr inp strong i j ++
          bulgeLIdx trnr inp strong i j ++
          bulgeRIdx trnr inp strong i j ++
          interior1xnLIdx trnr inp strong i j ++
          interior1xnRIdx trnr inp strong i j
          )
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z) <- str k l d'
      ] ++
      --
      -- multibranch loops closed at (i,j)
      --
      [ ((i,j):x++y,z)
      | i<j
      , (k,enext) <- multibranchCloseIdx trnr inp mbr mbr1 i j
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z') <- mul (i+1) k d'
      , (y,z) <- mul1 (k+1) (j-1) z' -- z' is what 'mul' left us
      ]
    mul i j d = let ehere = mbr!(i,j) in
      --
      -- unpaired nucleotide on the J side
      --
      [ (x,z)
      | i<j
      , ((k,l),enext) <- multibranchUnpairedJIdx trnr inp mbr i j
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z) <- mul k l d'
      ] ++
      --
      -- a helix (k,j)
      --
      [ (x,z)
      | i<j
      , (k,enext) <- multibranchKJHelixIdx trnr inp strong i j
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z) <- str k j d'
      ] ++
      --
      -- add a helix to an existing structure
      --
      [ (x++y,z)
      | i<j
      , (k,enext) <- multibranchAddKJHelixIdx trnr inp mbr strong i j
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z') <- mul i k d'
      , (y,z) <- str (k+1) j z'
      ]
    mul1 i j d = let ehere = mbr1!(i,j) in
      --
      -- Simply the strong part in a multibranch loop
      --
      [ (x,z)
      | i<j
      , (_,enext) <- multibranchIJLoopIdx trnr inp strong i j
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z) <- str i j d'
      ] ++
      --
      -- unpaired nucleotide on the J side
      --
      [ (x,z)
      | i<j
      , ((k,l),enext) <- multibranchUnpairedJIdx trnr inp mbr1 i j
      , let d' = newD d ehere enext
      , d'>=0
      , (x,z) <- mul1 k l d'
      ]

{-
test inp' k =
  let
    (_,n) = bounds inp
    bts = backtrack trnr inp tbls k
    inp = mkPrimary inp'
    trnr = fst turner2004GH
    tbls@(weak,strong,mbr1,mbr,extern) = fold trnr inp
    optE = extern ! (0,n)
    f x
      | x > 999 = "   x"
      | x < -999 = "-999"
      | otherwise = printf "%4d" x
    g t = do
            let ls = splitEvery (n+1) $ map f $ toList t
            putStr "    "
            mapM_ (printf "%4d") [0 :: Int .. (length $ head ls) -1]
            putStrLn ""
            zipWithM_ (\k xs -> printf "%4d" k >> mapM_ putStr xs >> putStrLn"") [0 :: Int ..] ls
  in do
      print "weak"
      g weak
      print "strong"
      g strong
      print "mbr   L"
      g mbr
      print "mbr1  R"
      g mbr1
      print "extern"
      g extern
      print optE
      putStrLn inp'
      mapM_ (\(s,e) -> (putStr $ dotbracket s) >> (putStrLn $ " " ++ show e)) bts
-}