import Control.Monad
import Control.Monad.Random

import Data.Bits
import Data.List

type Nucleotido = Bool

type Individuo = [Nucleotido]

type Poblacion = [Individuo]

type Probabilidad = Double

type TasaMutacion = Double

type ModeloAptitud = Individuo -> Double

type FuncionAjuste = Individuo -> Individuo

seleccion :: ModeloAptitud -> [(Individuo, Probabilidad)] -> [Individuo]
seleccion aptitud pob' = map (muestra . snd) pob'
  where
    pob = map fst pob'
    aptitudes = map aptitud pob
    totalAptitud = sum aptitudes
    prob = map (/ totalAptitud) aptitudes
    probAcum = tail (scanl (+) 0 prob)
    probAcum' = map (1-) probAcum
    elementos = zip probAcum' pob
    muestra u = snd . head . dropWhile (\x -> fst x > u) $ elementos

mutacion :: TasaMutacion -> [(Nucleotido, Probabilidad)] -> Individuo
mutacion pmut crom' = map mutar crom'
  where
    crom = map fst crom'
    mutar (b, u)
      | u < pmut  = not b
      | otherwise = b

cruce :: Probabilidad -> (Individuo, Individuo) -> (Individuo, Individuo)
cruce u (ind1, ind2) = (ind1inf ++ ind2sup, ind2inf ++ ind1sup)
  where
    long = fromIntegral (length ind1)
    puntoCruz = fromIntegral $ floor (u * long)
    (ind1inf, ind1sup) = splitAt puntoCruz ind1
    (ind2inf, ind2sup) = splitAt puntoCruz ind2

data ProblemaGenetico = PG  { pgAptitud :: ModeloAptitud
                            , pgAjuste  :: FuncionAjuste
                            , pgTasaMut :: TasaMutacion
                            , pgTamPob  :: Int
                            , pgNumGen  :: Int
                            , pgTamCrom :: Int
                            }

generacion :: (RandomGen g) => ProblemaGenetico -> Poblacion -> Rand g Poblacion
generacion pg pob = (ajustar >=> seleccionar) pob
  where
    ajustar pob = do
      return $ map (pgAjuste pg) pob
    
    seleccionar pob = do
      psel <- replicateM (pgTamPob pg) getRandom
      return $ seleccion (pgAptitud pg) (zip pob psel)
    
    cruzar pob = return pob
    
    mutar pob = forM pob $ \ind -> do
      pmut <- replicateM (pgTamPob pg) getRandom
      return $ mutacion (pgTasaMut pg) (zip ind pmut)

simulacion :: (RandomGen g) => ProblemaGenetico -> Rand g Poblacion
simulacion pg = do
  inicial <- replicateM (pgTamPob pg) genIndividuo
  foldM paso inicial [1..pgNumGen pg]
  where
    paso pob _ = generacion pg pob
    genIndividuo = replicateM (pgTamCrom pg) getRandom

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

logBase2 x = finiteBitSize x - 1 - countLeadingZeros x

data Planta = Planta Int Int Double Double Double

plantaLimite (Planta _ m _ _ _) = m

plantaLong (Planta n _ _ _ _) = n

plantaCosto (Planta _ _ a b c) x = a * x * x + b * x + c

problemaPlantas :: ProblemaGenetico
problemaPlantas = PG aptitud ajuste tasaMut tamPob numGen tamCrom

codificar :: [Int] -> Individuo
codificar sols = concat (zipWith cod plantas sols)
  where
    cod planta sol = map (testBit sol) [0..(plantaLong planta - 1)]

decodificar :: Individuo -> [Int]
decodificar ind = map dec partes
  where
    longitudes = map plantaLong plantas
    indices = init $ scanl (+) 0 longitudes
    partes = zipWith (\i n -> take n $ drop i ind) indices longitudes
    dec bs = foldl set zeroBits (zip bs [0..])
      where
        set x (b, i) = if b then x `setBit` i else x

costo :: Individuo -> Double
costo ind = sum $ zipWith plantaCosto plantas (map fromIntegral $ decodificar ind)

aptitud :: ModeloAptitud
aptitud ind = if falla == maxEnergia then 0 else 1 / (costo ind * (falla + 1))
  where
    sols = map fromIntegral $ decodificar ind
    falla = abs (maxEnergia - sum sols)

ajuste :: FuncionAjuste
ajuste ind = codificar sols'
  where
    sols  = decodificar ind
    sols' = zipWith ajustar plantas sols
    ajustar planta sol = sol `rem` (1 + plantaLimite planta)

tasaMut = 0.01
tamPob  = 1000
numGen  = 20
tamCrom = sum (map plantaLong plantas)

maxEnergia = 180

plantas = [ planta 100 0.3    (-18.0)   738.0
          , planta 100 0.166  (-6.66)   616.0
          , planta 100 0.208  (-12.08)  733.3
          , planta 100 0.255  (-15.4)   685.0 ]
  where
    planta m = Planta (1 + logBase2 m) m

main = do
  res <- evalRandIO (simulacion problemaPlantas)
  let mejor = last $ sortOn (pgAptitud problemaPlantas) res
  
  print (decodificar mejor)
  print (costo mejor)