{-# LANGUAGE Rank2Types #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Sequence.IonSeries
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Functions to work with a series of ions, and generate theoretical spectra
-- suitable for matching against experimental results with the sequest
-- cross-correlation algorithm.
--
--------------------------------------------------------------------------------

module Sequence.IonSeries
  (
    PeakSpectrum,
    extractPeaks,
    matchIons,
    matchIonSequence
  )
  where

import Mass
import Config
import Sequence.Fragment
import Spectrum.Data
import Spectrum.Correlation

import Control.Arrow
import Control.Monad.ST
import Data.Function
import Data.Vector.Algorithms.Intro
import qualified Data.ByteString.Lazy.Char8     as L
import qualified Data.Vector.Generic            as G
import qualified Data.Vector.Generic.Mutable    as M


type PeakSpectrum = Spectrum

--
-- Convert mass to mass/charge ratio
--
ionMZ :: Float -> Float -> Float
ionMZ m c = (m + massH*c) / c


--
-- Add a spectral peak for each fragment location, as well as the peaks
-- corresponding to neutral losses of H2O and NH3.
--
-- The factors that contribute to the collision induced dissociation process
-- used in tandem-MS experiments are not completely understood, so accurate
-- prediction of fragment ion abundances are not possible. Magnitude components
-- are assigned based on empirical knowledge.
--
addIonsAB, addIonsY :: Float -> Float -> [(Float, Float)]
addIonsAB charge mass = addIonsA : addIonsB
  where
    addIonsA = let m = ionMZ (mass + massH - massCO) charge in (m, 10)
    addIonsB = let m = ionMZ (mass + massH) charge in
      [
        (m,50), (m+1,25), (m-1,25),
        (m - massH2O/charge, 10),
        (m - massNH3/charge, 10)
      ]

addIonsY charge mass =
  let m = ionMZ (mass + massH + massOH) charge in
    [
      (m,50), (m+1,25), (m-1,25),
      (m - massNH3/charge, 10)
    ]


--------------------------------------------------------------------------------
-- Ion Matching
--------------------------------------------------------------------------------

--
-- The mass ladder generated by successive breaks of the peptide backbone from
-- the C- to N-terminus. This also includes the last element, corresponding to
-- the mass of the unbroken peptide, so that the y-ion ladder can be generated
-- from this.
--
-- This does not include neutral loss peaks.
--
bions :: ConfigParams -> Float -> L.ByteString -> [Float]
bions cp charge s = concatMap f [1 .. max 1 (charge-1)]
  where
    f c = map (`ionMZ` c) ms
    ms  = scanl1 (+) . map (getAAMass cp) . L.unpack $ s


--
-- Extract the 200 most intense peaks from an observed spectrum, collecting each
-- set of three adjacent bins into a single output.
--
{-# INLINE extractPeaks #-}
extractPeaks :: Spectrum -> PeakSpectrum
extractPeaks spec = peaks
  where
    zeros = G.replicate ((G.length spec `div` 3) + 1) 0
    peaks = G.accumulate max zeros
          . G.take 200
          . apply (flip (selectBy (flip compare `on` snd)) 200)
          . G.imap (\i v -> (i `div` 3, v))
          $ spec

{-# INLINE apply #-}
apply :: G.Vector v e => (forall s mv. M.MVector mv e => mv s e -> ST s ()) -> v e -> v e
apply algo vec = G.create $ G.thaw vec >>= \mvec -> algo mvec >> return mvec


--
-- Calculate the number of predicted fragment ions which match the peak spectrum
-- to within 1 da. Returns the number matched together with the total number of
-- ions. This does not consider neutral loss.
--
{-# INLINE matchIons #-}
matchIons :: ConfigParams -> Float -> PeakSpectrum -> Fragment -> (Int, Int)
matchIons cp charge peaks = matched &&& total
  where
    boolToInt True  = 1
    boolToInt False = 0

    total   f = fromIntegral $ (L.length (fragdata f) - 5) * 2
    matched f = let (b,y) = matchIonSequence cp charge peaks f
                in  sum   $ map boolToInt (b ++ y)


--
-- Determine which of the predicted fragment ions match the observed peak
-- spectrum, for the b- and y-ion sequences respectively.
--
{-# INLINE matchIonSequence #-}
matchIonSequence :: ConfigParams -> Float -> PeakSpectrum -> Fragment -> ([Bool], [Bool])
matchIonSequence cp charge peaks (Fragment r _ f) =
  ( init (matched ions)
  , tail (matched (map (r-) ions))
  )
  where
    bin x = round (x / binWidth cp / 3)
    pep   = L.drop 2 $ L.take (L.length f - 2) f
    ions  = bions cp charge pep

    inrange i = i >= 0 && i < G.length peaks    -- unnecessary?
    matched   = map ((\i -> inrange i && peaks G.! i > 0) . bin)

