--------------------------------------------------------------------------------
-- |
-- Module      :  HarmTrace.Accompany
-- Copyright   :  (c) 2010-2012 Universiteit Utrecht, 2012 University of Oxford
-- License     :  GPL3
--
-- Maintainer  :  bash@cs.uu.nl, jpm@cs.ox.ac.uk
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Summary: Generate a melody (accompaniment) for a given harmony
--------------------------------------------------------------------------------

module HarmTrace.Accompany where

import HarmTrace.Base.MusicRep
import HarmTrace.Song
import HarmTrace.Models.Simple.Model    ( Piece )
import HarmTrace.Models.Simple.Main     ( getChords )
import HarmTrace.Models.ChordTokens     ( ChordToken(..), ctToCL, sdToNote )

import Control.Monad.State
import System.Random

import Data.List                        ( intersect )

import Debug.Trace


map2 :: (b -> c) -> [(a,b)] -> [(a,c)]
map2 f = map (\(a,b) -> (a, f b))

data MyState = MyState { genState     :: StdGen
                       , keyState     :: Key
                       , pieceState   :: Piece
                       , chordsState  :: [ChordToken] }

accompanyIO :: Key -> Piece -> IO Song
accompanyIO k p = do gen <- getStdGen
                     let initState = MyState gen k p (getChords p)
                     return (evalState (accompany k) initState)

accompany :: Key -> State MyState Song
accompany k =     allPossible >>= refine >>= pickBest>>= embellish
              >>= return . Song k

-- 2.1) Generate candidate melody notes from chords
allPossibleRel :: [ChordToken] -> [(ChordToken,[ScaleDegree])]
allPossibleRel cs = [ (c, notesRootedOn (root c) (classType c)) | c <- cs ]

allPossible :: State MyState [(ChordLabel,[MelodyNote])]
allPossible = do k <- get >>= return . keyState
                 p <- get >>= return . pieceState
                 let unRel (c,sds) = ( ctToCL k c
                                     , map (flip MelodyNote 3 . sdToNote k) sds)
                 return $ map unRel (allPossibleRel (getChords p))

-- All notes that belong to this chord
notesRootedOn :: ScaleDegree -> ClassType -> [ScaleDegree]
notesRootedOn sd cls =
  let indices = case cls of
                  MajClass -> [0,4,7]
                  MinClass -> [0,3,7]
                  DimClass -> [0,3,6]
                  DomClass -> [0,4,7,10]
                  x        -> error $ "notesRootedOn: " ++ show x
  in [ transposeSem sd i | i <- indices ]

-- 2.2) Trim to remove bad candidates
refine :: [(ChordLabel, [MelodyNote])] -> State MyState [(ChordLabel, [MelodyNote])]
refine []            = return []
refine ((cl,mns):cs) = 
  do k <- get >>= return . keyState
     -- Make sure the first note is one of I, III, or V
     let indices = case keyMode k of
                     MajMode -> [0,4,7::Int]
                     MinMode -> [0,3,7]
         ki = toSemitone (keyRoot k)
         makeNote i = MelodyNote (toRootM (i + ki)) 3
         first = map makeNote indices
         firstNotes = let wanted = first `intersect` mns
                      in if null wanted then mns else wanted
         -- Handle the final note
         lastNote ns = let (a,[b]) = splitAt (length ns - 1) ns
                       in a ++ [final b]
         -- We want the final note to be a I, or, if I is not in the chord, a V
         final (c,n) = let n' = if makeNote 0 `elem` n
                                then [makeNote 0]
                                else [makeNote 7]
                       in (c,n')
     return $ ((cl,firstNotes) : lastNote cs)

-- 2.3) Pick one note per chord
pickBest :: [(ChordLabel, [MelodyNote])] -> State MyState [(ChordLabel, MelodyNote)]
pickBest cs =
  do s <- get
     -- Just randomly pick notes from the list of candidates
     let g = genState s
         rs = randoms g
         f ((cl, mns), r) = (cl, mns !! (r `mod` length mns))
         result = map f (zip cs rs)
         -- Make sure VIIs are followed by a I in the right octave
         k  = keyState s
         ki = toSemitone (keyRoot k)
         makeNote i = MelodyNote (toRootM (i + ki)) 3
         resolveCadences :: [(ChordLabel, MelodyNote)] -> [(ChordLabel, MelodyNote)]
         resolveCadences ((c1,n1):(c2,n2):cns)
           | n1 == makeNote 0 && n2 == makeNote 11
           = (c1,n1) : (c2,octaveDown n2) : resolveCadences cns
           | n1 == makeNote 11 && n2 == makeNote 0
           = (c1,n1) : (c2,octaveUp n2) : resolveCadences cns
           | otherwise = (c1,n1) : resolveCadences ((c2,n2):cns)
         resolveCadences x = x
     return (resolveCadences result)

-- 2.4) Embellish the melody
embellish :: [(ChordLabel, MelodyNote)] -> State MyState [(ChordLabel, [MelodyNote])]
embellish []            = return []
embellish ((cl,mn):cls) = do g <- get >>= return . genState
                             k <- get >>= return . keyState
                             return $ go k (cl,mn,g) cls
  where
    go k (cl1,n1,g1) []             = [(cl1,[n1])]
    go k (cl1,n1,g1) ((cl2,n2):cls) = let (_,g2) = next g1
                                      in (cl1, connectNotes g1 k cl1 n1 n2)
                                           : go k (cl2,n2,g2) cls

-- Given two notes, return a melody that begins in the first and ends in
-- something suitable to be connected to the second.
connectNotes :: StdGen -> Key -> ChordLabel
             -> MelodyNote -> MelodyNote -> [MelodyNote]
connectNotes g k cl n1@(MelodyNote r1 o1) n2@(MelodyNote r2 o2)
  | n1 == n2
-- embellish repetitions
  = let scale = notesInChord cl -- this is debatable...
    in case fst (randomR (0,3::Int) g) of
         -- C C -> C C
         0 -> [n1]
         -- C C -> C D E C
         1 -> if n1 `elem` scale 
              then take 3 . dropWhile (/= n1) $ scale
              else [n1]
         -- C C -> C B C
         2 -> if n1 `elem` scale 
              then take 2 . dropWhile (/= n1) $ reverse scale
              else [n1]
         -- C C -> C E D C
         3 -> let f123_132 [c,d,e] = [c,e,d]
              in if n1 `elem` scale 
                 then f123_132 . take 3 . dropWhile (/= n1) $ scale
                 else [n1]
         _ -> error "connectNotes: impossible"

-- Connect from a scale in the current key
connectNotes g k cl n1@(MelodyNote r1 o1) n2@(MelodyNote r2 o2)
  = let scale = notesInKey k
        -- scale = notesInChord cl
        line = if n1 < n2
               then n1 : takeWhile (< n2) (dropWhile (<= n1)          scale)
               else n1 : takeWhile (n2 <) (dropWhile (n1 <=) (reverse scale))
    in line

-- Like toRoot, but mod 12
toRootM :: Int -> Root
toRootM = toRoot . (`mod` 12)

-- Returns a scale
notesInKey :: Key -> [MelodyNote]
notesInKey (Key r m) = let indices = case m of
                                       MajMode -> [0,2,4,5,7,9,11]
                                       MinMode -> [0,2,3,5,7,8,10] -- tricky
                           base = [ toRootM (toSemitone r + i) | i <- indices ]
                       in filter (\n -> mnRoot n `elem` base) allMelodyNotes

-- Return the scale associated with a chord
notesInChord :: ChordLabel -> [MelodyNote]
notesInChord cl = let r = chordRoot cl
                      indices = case chordShorthand cl of
                                  Maj -> [0,2,4,5,7,9,11]
                                  Sev -> [0,2,4,5,7,9,10]
                                  Min -> [0,2,3,5,7,8,10]
                                  Dim -> [0,2,3,5,6,8,9]
                                  m   -> error $ "notesInChord: " ++ show m
                      base = [ toRootM (toSemitone r + i) | i <- indices ]
                  in filter (\n -> mnRoot n `elem` base) allMelodyNotes