package main

import (
        "flag"
        "fmt"

        "github.com/Kilemonn/MarkovChainSimulator/config"
)

const (
        OPTION_FILE_PATH = "f"
)

func main() {
        var configFilePath string

        flag.StringVar(&configFilePath, OPTION_FILE_PATH, "", "configuration file path")
        flag.Parse()

        if configFilePath == "" {
                fmt.Printf("No file path provided with [-%s] argument option\n", OPTION_FILE_PATH)
                return
        }

        chain, err := config.FromYaml(configFilePath)
        if err != nil {
                fmt.Printf("Failed to read in yaml file at path [%s]. Error: %s\n", configFilePath, err.Error())
                return
        }

        err = chain.Validate()
        if err != nil {
                fmt.Printf("Configuration validation failed with error: %s\n", err.Error())
                return
        }

        result := chain.Simulate()
        for k, v := range result {
                fmt.Printf("%s - %d/%d = %.2f%%\n", k, v, chain.StepCount, (float64(v)/float64(chain.StepCount))*100)
        }
}

package config

import (
        "fmt"
        "math/rand/v2"
        "os"

        "gopkg.in/yaml.v3"
)

type MarkovChainConfig struct {
        InitialState string `yaml:"initialState"`
        StepCount    uint   `yaml:"stepCount"`
        States       []MarkovChainState
}

type MarkovChainState struct {
        Id          string
        Transitions []MarkovChainTransition
}

type MarkovChainTransition struct {
        State  string
        Chance float64
}

func (config MarkovChainConfig) Validate() error {
        if len(config.States) == 0 {
                return fmt.Errorf("no states provided in configuration")
        }

        if !config.hasState(config.InitialState) {
                return fmt.Errorf("initial state [%s] does not exist in states list", config.InitialState)
        }

        for _, state := range config.States {
                chanceTotal := 0.0
                if len(state.Transitions) == 0 {
                        return fmt.Errorf("no transitions from state [%s] are defined", state.Id)
                }

                for _, transition := range state.Transitions {
                        if !config.hasState(transition.State) {
                                return fmt.Errorf("transition to state [%s] defined in transition from state [%s] does not exist",
                                        transition.State, state.Id)
                        }

                        if transition.Chance <= 0 {
                                return fmt.Errorf("undefined or non-positive chance value [%f] set for transition to state [%s] from state [%s]",
                                        transition.Chance, transition.State, state.Id)
                        }
                        chanceTotal += transition.Chance
                }

                if chanceTotal != 1.0 {
                        return fmt.Errorf("combined chance total does not equal 1.0 but adds up to [%f] for transitions defined for state [%s]",
                                chanceTotal, state.Id)
                }
        }

        return nil
}

func (config MarkovChainConfig) hasState(stateId string) bool {
        if len(config.States) == 0 {
                return false
        }

        for _, state := range config.States {
                if state.Id == stateId {
                        return true
                }
        }

        return false
}

func (config MarkovChainConfig) Simulate() map[string]int {
        state := config.getStateById(config.InitialState)
        var count uint = 0
        stateCount := make(map[string]int)

        for count < config.StepCount {
                if val, exists := stateCount[state.Id]; exists {
                        stateCount[state.Id] = val + 1
                } else {
                        stateCount[state.Id] = 1
                }

                state = state.getNextState(config)

                count += 1
        }

        return stateCount
}

func (config MarkovChainConfig) getStateById(stateId string) MarkovChainState {
        for _, state := range config.States {
                if state.Id == stateId {
                        return state
                }
        }

        return MarkovChainState{}
}

func (state MarkovChainState) getNextState(config MarkovChainConfig) MarkovChainState {
        rng := rand.Float64()

        val := 0.0
        lastStateName := ""

        for _, transition := range state.Transitions {
                val += transition.Chance
                if rng <= val {
                        return config.getStateById(transition.State)
                }
                lastStateName = transition.State
        }

        return config.getStateById(lastStateName)
}

func FromYaml(configFilePath string) (config MarkovChainConfig, err error) {
        data, err := os.ReadFile(configFilePath)
        if err != nil {
                fmt.Printf("Failed to read in file [%s]. Error: [%s].\n", configFilePath, err.Error())
                return
        }

        err = yaml.Unmarshal(data, &config)
        if err != nil {
                fmt.Printf("Failed to unmarshal data from file [%s]. Error: [%s].\n", configFilePath, err.Error())
        }

        return
}