d3-worldcup/teams.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}

import           Data.Aeson                 as AE
import           Data.Aeson.Types           as AET
import           Data.ByteString.Lazy       as BL
import           Data.ByteString.Lazy.Char8 as BL8
import           Data.HashMap.Strict        as HM
import           Data.Maybe                 as M
import           Prelude                    as P

----- Data constructors -----
data TeamFinal = TeamFinal {
    resultId        :: Int,
    resultCountry   :: String,
    resultContinent :: String,
    resultGoals     :: Int
} deriving (Show)

data Team = Team {
    teamId        :: Int,
    teamCountryId :: Int } deriving (Show)

data EventTeam = EventTeam {
    eventTeamId   :: Int } deriving (Show)

data Country = Country {
    countryId          :: Int,
    countryName        :: String,
    countryContinentId :: Int } deriving (Show)

data Continent = Continent {
    continentId   :: Int,
    continentName :: String } deriving (Show)

data Game = Game {
    roundId  :: Int,
    playAt   :: String,
    team1Id  :: Int,
    team2Id  :: Int,
    score1   :: Value,
    score2   :: Value,
    score1et :: Value,
    score2et :: Value,
    score1p  :: Value,
    score2p  :: Value } deriving (Show)

----- Type synonyms -----
type Teams = HashMap Int Team
type Countries = HashMap Int Country
type Continents = HashMap Int Continent

----- ToJSON parsers -----
instance ToJSON TeamFinal where
    toJSON TeamFinal{..} =
        object [ "id" .= resultId, "ct" .= resultCountry, "cn" .= resultContinent, "g" .= resultGoals ]

----- FromJSON parsers -----
instance FromJSON Team where
    parseJSON = withObject "team" $ \o -> do
        teamId <- o .: "id"
        teamCountryId <- o .: "country_id"
        return Team{..}

instance FromJSON EventTeam where
    parseJSON = withObject "team" $ \o -> do
        eventTeamId <- o .: "team_id"
        return EventTeam{..}

instance FromJSON Game where
    parseJSON = withObject "game" $ \o -> do
        roundId <- o .: "round_id"
        team1Id <- o .: "team1_id"
        team2Id <- o .: "team2_id"
        playAt <- o .: "play_at"
        score1 <- o .: "score1"
        score2 <- o .: "score2"
        score1et <- o .: "score1et"
        score2et <- o .: "score2et"
        score1p <- o .: "score1p"
        score2p <- o .: "score2p"
        return Game{..}

instance FromJSON Country where
    parseJSON = withObject "country" $ \o -> do
        countryId <- o .: "id"
        countryName <- o .: "name"
        countryContinentId <- o .: "continent_id"
        return Country{..}

instance FromJSON Continent where
    parseJSON = withObject "continent" $ \o -> do
        continentId <- o .: "id"
        continentName <- o .: "name"
        return Continent{..}

parseTeams :: Either String [Value] -> Teams
parseTeams (Left x)   = error x
parseTeams (Right xs) = P.foldl reduce HM.empty xs where
    reduce acc x = case (parseEither parseJSON x :: Either String Team) of
        (Left s)  -> error s
        (Right v) -> HM.insert (teamId v) v acc

parseEventTeams :: Either String [Value] -> [Int]
parseEventTeams (Left x)   = error x
parseEventTeams (Right xs) = P.foldl reduce [] xs where
    reduce acc x = case (parseEither parseJSON x :: Either String EventTeam) of
        (Left s)  -> error s
        (Right v) -> eventTeamId v : acc

parseGames :: Either String [Value] -> [Game]
parseGames (Left x)   = error x
parseGames (Right xs) = P.foldl reduce [] xs where
    reduce acc x = case (parseEither parseJSON x :: Either String Game) of
        (Left s)  -> error s
        (Right v) -> v : acc

parseCountries :: Either String [Value] -> Countries
parseCountries (Left x)   = error x
parseCountries (Right xs) = P.foldl reduce HM.empty xs where
    reduce acc x = case (parseEither parseJSON x :: Either String Country) of
        (Left s)  -> error s
        (Right v) -> HM.insert (countryId v) v acc

parseContinents :: Either String [Value] -> Continents
parseContinents (Left x)   = error x
parseContinents (Right xs) = P.foldl reduce HM.empty xs where
    reduce acc x = case (parseEither parseJSON x :: Either String Continent) of
        (Left s)  -> error s
        (Right v) -> HM.insert (continentId v) v acc

parseInt :: Value -> Int
parseInt Null = 0
parseInt x    = read $ BL8.unpack $ encode x :: Int

----- Hashmap lookup convenience -----
lookupTeam :: Int -> Teams -> Team
lookupTeam _id hm = M.fromMaybe (Team 999 999) (HM.lookup _id hm)

lookupCountry :: Int -> Countries -> Country
lookupCountry _id hm = M.fromMaybe (Country 999 "Unknown" 999) (HM.lookup _id hm)

lookupContinent :: Int -> Continents -> Continent
lookupContinent _id hm = M.fromMaybe (Continent 999 "Unknown") (HM.lookup _id hm)

countryFromTeamId :: Int -> Teams -> Countries -> Country
countryFromTeamId _id hmT = lookupCountry a where
    a = teamCountryId $ lookupTeam _id hmT

continentFromTeamId :: Int -> Teams -> Countries -> Continents -> Continent
continentFromTeamId _id hmT hmC = lookupContinent (countryContinentId b) where
    a = teamCountryId $ lookupTeam _id hmT
    b = lookupCountry a hmC

goalsFromTeamId :: Int -> [Game] -> Int
goalsFromTeamId _id = P.foldl reducer 0 where
    reducer acc x
        | team1Id x == _id = acc + parseInt (score1 x) + parseInt (score1et x) + parseInt (score1p x)
        | team2Id x == _id = acc + parseInt (score2 x) + parseInt (score2et x) + parseInt (score2p x)
        | otherwise = acc

buildResult :: [Int] -> Teams -> Countries -> Continents -> [Game] -> [TeamFinal]
buildResult ets hmT hmC hmN gs = P.foldl reducer [] ets where
    reducer acc x = TeamFinal x (getCountryName x) (getContinentName x) (goalCount x) : acc
    getCountryName x = countryName (countryFromTeamId x hmT hmC)
    getContinentName x = continentName (continentFromTeamId x hmT hmC hmN)
    goalCount x = goalsFromTeamId x gs

main :: IO ()
main = do
    dataTeams <- BL.readFile "./data/teams.json"
    dataEventTeams <- BL.readFile "./data/events_teams.json"
    dataCountries <- BL.readFile "./data/countries.json"
    dataContinents <- BL.readFile "./data/continents.json"
    dataGames <- BL.readFile "./data/games.json"

    let teams = parseTeams (AE.eitherDecode dataTeams)
    let eventteams = parseEventTeams (AE.eitherDecode dataEventTeams)
    let games = parseGames (AE.eitherDecode dataGames)
    let countries = parseCountries (AE.eitherDecode dataCountries)
    let continents = parseContinents (AE.eitherDecode dataContinents)

    let encoded = encode $ buildResult eventteams teams countries continents games
    BL8.putStrLn encoded
    P.putStrLn $ show (P.length eventteams) ++ " teams found."
    BL8.putStrLn "Writing teams.json"
    BL8.writeFile "teams.json" encoded