module BioInf.RNAfold.Energy where
import Control.Exception (assert)
import Data.Strict.Tuple hiding (fst,snd)
import qualified Data.Vector.Fusion.Stream.Monadic as S
import Biobase.Primary
import Biobase.Secondary.Vienna
import Biobase.Vienna
import Data.PrimitiveArray
import "PrimitiveArray" Data.Array.Repa.Index
import Debug.Trace
hairpinF :: Vienna2004 -> (Nuc :!: Nuc) -> (Primary :!: Int :!: Int) -> (Nuc :!: Nuc) -> Int
hairpinF Vienna2004{..} (l :!: lp) (xs :!: i :!: j) (rp :!: r) = hairpinType where
hairpinType
| len < 3 = 999999
| len == 3 = hairpinL!(Z:.len) + tAU
| len > 31 = hairpinL!(Z:.30) + hairpinMM!(Z:.p:.lp:.rp) + llp
| otherwise = hairpinL!(Z:.len) + hairpinMM!(Z:.p:.lp:.rp)
p = mkViennaPair (l,r)
len = ji+1
tAU = if p/=vpCG && p/=vpGC then termAU else 0
llp = floor $ 108.856 * log (fromIntegral len / 30)
tinyloopF :: Vienna2004 -> (Primary :!: Int :!: Int) -> Int -> (Primary :!: Int :!: Int) -> Int
tinyloopF Vienna2004{..} (ls :!: li :!: lj) w (rs :!: ri :!: rj) = assert (assertL && assertR) $ loopType where
assertL = let (_,Z:.n) = bounds ls in li>=0 && lj<=n
assertR = let (_,Z:.n) = bounds rs in ri>=0 && rj<=n
loopType
| dl==0 || dr==0 = error $ "bug in tinyloop: " ++ show (li,lj,[ls!(Z:.k) | k<-[li..lj]],ri,rj,[rs!(Z:.l) | l<-[ri..rj]])
| dl==1 && dr==1 = w+stack!(Z:.op:.ip)
| dl==1 && dr==2 = w+stack!(Z:.op:.ip)+bulgeL!(Z:.1)
| dl==2 && dr==1 = w+stack!(Z:.op:.ip)+bulgeL!(Z:.1)
| dl==2 && dr==2 = w+iloop1x1!(Z:.op:.ip:.bI:.bJ)
| dl==2 && dr==3 = w+iloop2x1!(Z:.op:.ip:.bI:.bL:.bJ)
| dl==3 && dr==2 = w+iloop2x1!(Z:.op:.ip:.bJ:.bI:.bK)
| dl==3 && dr==3 = w+iloop2x2!(Z:.op:.ip:.bI:.bK:.bL:.bJ)
| dl==3 && dr==4
|| dl==4 && dr==3 = w+iloop2x3MM!(Z:.op:.bI:.bJ) + iloop2x3MM!(Z:.ip:.bL:.bK) + iloopL!(Z:.5) + ninio
| otherwise = 999999
op = mkViennaPair (ls!(Z:.li),rs!(Z:.rj))
ip = mkViennaPair (rs!(Z:.ri),ls!(Z:.lj))
bI = ls!(Z:.li+1)
bJ = rs!(Z:.rj1)
bK = ls!(Z:.lj1)
bL = rs!(Z:.ri+1)
dl = ljli
dr = rjri
bulgeLF :: Vienna2004 -> (Primary :!: Int :!: Int) -> Int -> (Nuc :!: Nuc) -> Int
bulgeLF Vienna2004{..} (ls :!: li :!: lj) w (rp :!: r) = assert (ljli<=30) $ w + tAUlr + lenE + tAUrpcp where
tAUlr = terminalAU termAU lr
tAUrpcp = terminalAU termAU rpcp
lr = mkViennaPair (ls!(Z:.li),r)
rpcp = mkViennaPair (rp,ls!(Z:.lj))
lenE = bulgeL!(Z:.ljli1)
bulgeRF :: Vienna2004 -> (Nuc :!: Nuc) -> Int -> (Primary :!: Int :!: Int) -> Int
bulgeRF Vienna2004{..} (l :!: lp) w (rs :!: ri :!: rj) = assert (rjri<=30) $ w + tAUlr + lenE + tAUcplp where
tAUlr = terminalAU termAU lr
tAUcplp = terminalAU termAU cplp
lr = mkViennaPair (l,rs!(Z:.rj))
cplp = mkViennaPair (rs!(Z:.ri),lp)
lenE = bulgeL!(Z:.rjri1)
iloopN1F :: Vienna2004 -> (Primary :!: Int :!: Int) -> Int -> (Primary :!: Int :!: Int) -> Int
iloopN1F Vienna2004{..} (ls:!:li:!:lj) w (rs:!:ri:!:rj)
= assert (ljli1 <=29 && rjri == 2)
$ w + outerMM + lenE + innerMM + owninio where
poLR = mkViennaPair (ls!(Z:.li), rs!(Z:.rj))
piRL = mkViennaPair (rs!(Z:.ri), ls!(Z:.lj))
outerMM = iloop1xnMM!(Z:.poLR:.ls!(Z:.li+1):.rs!(Z:.ri+1))
innerMM = iloop1xnMM!(Z:.piRL:.rs!(Z:.ri+1):.ls!(Z:.lj1))
lenE = iloopL!(Z:.(lL+1))
owninio = min maxNinio (ninio * (lL1))
lL = ljli1
lR = rjri1
iloop1NF :: Vienna2004 -> (Primary :!: Int :!: Int) -> Int -> (Primary :!: Int :!: Int) -> Int
iloop1NF Vienna2004{..} (ls :!: li :!: lj) w (rs :!: ri :!: rj)
= assert (ljli == 2 && rjri1 <=29)
$ w + outerMM + lenE + innerMM + owninio where
poLR = mkViennaPair (ls!(Z:.li), rs!(Z:.rj))
piRL = mkViennaPair (rs!(Z:.ri), ls!(Z:.lj))
outerMM = iloop1xnMM!(Z:.poLR:.ls!(Z:.li+1):.rs!(Z:.rj1))
innerMM = iloop1xnMM!(Z:.piRL:.rs!(Z:.ri+1):.ls!(Z:.li+1))
lenE = iloopL!(Z:.(lR+1))
owninio = min maxNinio (ninio * (lR1))
lL = ljli1
lR = rjri1
iloopIF :: Vienna2004 -> (Primary :!: Int :!: Int) -> Int -> (Primary :!: Int :!: Int) -> Int
iloopIF Vienna2004{..} (ls :!: li :!: lj) w (rs :!: ri :!: rj) = w + iloopMM!(Z:.p:.bR:.bL) + iloopL!(Z:.len) + owninio where
p = mkViennaPair (rs!(Z:.ri), ls!(Z:.lj))
bL = ls!(Z:.lj1)
bR = rs!(Z:.ri+1)
len = (ljli)+(rjri)
owninio = min maxNinio (ninio * (abs $ (ljli) (rjri)))
iloopOF :: Vienna2004 -> (Nuc :!: Nuc) -> Int -> (Nuc :!: Nuc) -> Int
iloopOF Vienna2004{..} (l :!: lp) iloopif (rp :!: r) = iloopif + iloopMM!(Z:.op:.lp:.rp)
where op = mkViennaPair (l,r)
multiIF :: Vienna2004 -> Int -> Int -> Int
multiIF Vienna2004{..} b c = b+c
multiOF :: Vienna2004 -> (Nuc :!: Nuc) -> Int -> (Nuc :!: Nuc) -> Int
multiOF Vienna2004{..} (l :!: lp) multiif (rp :!: r) = multiif + multiMM!(Z:.p:.lp:.rp) + multiHelix + multiOffset + terminalAU termAU p
where p = mkViennaPair (l,r)
regionStemF :: Vienna2004 -> Nuc -> Int -> Int
regionStemF Vienna2004{..} _ w = w + multiNuc
justStemF :: Vienna2004 -> (Nuc :!: Nuc) -> Int -> (Nuc :!: Nuc) -> Int
justStemF Vienna2004{..} (l :!: lp) w (rp :!: r) = w + multiMM!(Z:.p:.l:.r)
where p = mkViennaPair (lp,rp)
bsF :: Int -> Int -> Int
bsF b reg = b
rSF :: Nuc -> Int -> Int
rSF nuc w = w
bcF :: Int -> Int -> Int
bcF b c = b + c
ssF :: Int -> Int
ssF reg = 0
cmF :: Int -> Int -> Int
cmF w s = w+s
nilF :: Bool -> Int
nilF e = if e then 0 else 999999
iD = id
h :: Monad m => S.Stream m Int -> m Int
h = S.foldl' min (999999::Int)
terminalAU termAU p = if p/=vpCG && p/=vpGC then termAU else 0