#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Entry point for running the GABM application.
To run: python3 -m gabm
"""
# Metadata
__author__ = ["Andy Turner <agdturner@gmail.com>"]
__version__ = "0.4.0"
__copyright__ = "Copyright (c) 2026 GABM contributors, University of Leeds"

# Standard library imports
import argparse
import os
import sys
import logging
from pathlib import Path
import random
# Visualization
import matplotlib.pyplot as plt
# Local imports
from gabm.abm.environment import Environment, Nation
from gabm.abm.agent import Agent, Person
from gabm.abm.group import Group, OpinionatedGroup
from gabm.abm.attributes.opinion import OpinionTopicID, OpinionValue, OpinionValueMap, OpinionTopic, Opinion


def run_example():
    """
    Run the original example simulation.
    """
    # ...existing code...
    logging.info("\n--- GABM ---\n")
    # Set random seed for reproducibility
    random.seed(42)
    # Flexible group sizes
    n_negative = 2
    n_positive = 2
    n_neutral = 6
    # Number of communication rounds
    n_iterations = 5
    # For plotting: record opinions at each round (including initial)
    opinions_over_time = []
    # Initialize the environment
    env = Environment()
    # Define opinion topics and values
    # Define opinion topic IDs
    negative_opinion_topic_id = OpinionTopicID(0)
    neutral_opinion_topic_id = OpinionTopicID(1)
    positive_opinion_topic_id = OpinionTopicID(2)
    # Define opinion topics
    negative_opinion_topic = OpinionTopic(negative_opinion_topic_id, "negative", "A negative opinion.")
    neutral_opinion_topic = OpinionTopic(neutral_opinion_topic_id, "neutral", "A neutral opinion.")
    positive_opinion_topic = OpinionTopic(positive_opinion_topic_id, "positive", "A positive opinion.")
    # Add topics to environment
    env.opinions = {
        negative_opinion_topic_id: negative_opinion_topic,
        neutral_opinion_topic_id: neutral_opinion_topic,
        positive_opinion_topic_id: positive_opinion_topic
    }
    # Define opinion values for each topic
    env.opinion_values = {
        negative_opinion_topic_id: {
            OpinionValue(negative_opinion_topic_id, -2, "Strongly negative"),
            OpinionValue(negative_opinion_topic_id, -1, "Negative"),
            OpinionValue(negative_opinion_topic_id, 0, "Neutral"),
            OpinionValue(negative_opinion_topic_id, 1, "Positive"),
            OpinionValue(negative_opinion_topic_id, 2, "Strongly positive")
        },
        neutral_opinion_topic_id: {
            OpinionValue(neutral_opinion_topic_id, -2, "Strongly negative"),
            OpinionValue(neutral_opinion_topic_id, -1, "Negative"),
            OpinionValue(neutral_opinion_topic_id, 0, "Neutral"),
            OpinionValue(neutral_opinion_topic_id, 1, "Positive"),
            OpinionValue(neutral_opinion_topic_id, 2, "Strongly positive")
        },
        positive_opinion_topic_id: {
            OpinionValue(positive_opinion_topic_id, -2, "Strongly negative"),
            OpinionValue(positive_opinion_topic_id, -1, "Negative"),
            OpinionValue(positive_opinion_topic_id, 0, "Neutral"),
            OpinionValue(positive_opinion_topic_id, 1, "Positive"),
            OpinionValue(positive_opinion_topic_id, 2, "Strongly positive")
        }
    }
    # Define opinion values for easy access
    env.opinion_value_map = {}
    for topic_id, values in env.opinion_values.items():
        for value in values:
            env.opinion_value_map[value.opinion_topic_id] = value
    # Create negative agents
    negative = env.groups_active[0] = Group(0, "Negative")
    negative_opinions = {
        negative_opinion_topic_id: Opinion(
            negative_opinion_topic_id,
            OpinionValueMap({negative_opinion_topic_id: OpinionValue(negative_opinion_topic_id, 2, "Strongly positive")}),
            2),
        neutral_opinion_topic_id: Opinion(
            neutral_opinion_topic_id,
            OpinionValueMap({neutral_opinion_topic_id: OpinionValue(neutral_opinion_topic_id, 0, "Neutral")}),
            0),
        positive_opinion_topic_id: Opinion(
            positive_opinion_topic_id,
            OpinionValueMap({positive_opinion_topic_id: OpinionValue(positive_opinion_topic_id, -2, "Strongly negative")}),
            -2)
    }
    for agent_id in range(n_negative):
        env.agents_active[agent_id] = Person(
            agent_id, environment=env, year_of_birth=2000,
            opinions=negative_opinions)
        negative.add_member(env.agents_active[agent_id])
    # Create positive agents
    positive = env.groups_active[1] = Group(1, "Positive")
    positive_opinions = {
        negative_opinion_topic_id: Opinion(
            negative_opinion_topic_id,
            OpinionValueMap({negative_opinion_topic_id: OpinionValue(negative_opinion_topic_id, -2, "Strongly negative")}),
            -2),
        neutral_opinion_topic_id: Opinion(
            neutral_opinion_topic_id,
            OpinionValueMap({neutral_opinion_topic_id: OpinionValue(neutral_opinion_topic_id, 0, "Neutral")}),
            0),
        positive_opinion_topic_id: Opinion(
            positive_opinion_topic_id,
            OpinionValueMap({positive_opinion_topic_id: OpinionValue(positive_opinion_topic_id, 2, "Strongly positive")}),
            2)
    }
    for agent_id in range(n_negative, n_negative + n_positive):
        env.agents_active[agent_id] = Person(
            agent_id, environment=env, year_of_birth=2000, 
            opinions=positive_opinions)
        positive.add_member(env.agents_active[agent_id])
    # Create neutral agents
    neutral = env.groups_active[2] = Group(2, "Neutral")
    neutral_opinions = {
        negative_opinion_topic_id: Opinion(
            negative_opinion_topic_id,
            OpinionValueMap({negative_opinion_topic_id: OpinionValue(negative_opinion_topic_id, 0, "Neutral")}),
            0),
        neutral_opinion_topic_id: Opinion(
            neutral_opinion_topic_id,
            OpinionValueMap({neutral_opinion_topic_id: OpinionValue(neutral_opinion_topic_id, 2, "Strongly positive")}),
            2),
        positive_opinion_topic_id: Opinion(
            positive_opinion_topic_id,
            OpinionValueMap({positive_opinion_topic_id: OpinionValue(positive_opinion_topic_id, 0, "Neutral")}),
            0)
    }
    for agent_id in range(n_negative + n_positive, n_negative + n_positive + n_neutral):
        env.agents_active[agent_id] = Person(
            agent_id, environment=env, year_of_birth=2000,
            opinions=neutral_opinions)
        neutral.add_member(env.agents_active[agent_id])
    # Log the initial state of the environment
    n_agents = len(env.agents_active)
    logging.info(f"Initialized environment with {n_agents} agents.")
    # Record initial opinions
    opinions_over_time.append(
        [agent.opinions.copy() for agent in env.agents_active.values()]
    )
    # Calculate the average opinions of all agents in the environment and log it.
    topic_name_to_id = {
        "negative": negative_opinion_topic_id,
        "neutral": neutral_opinion_topic_id,
        "positive": positive_opinion_topic_id
    }
    for topic_name, topic_id in topic_name_to_id.items():
        avg_opinion = sum(
            agent.opinions[topic_id].value if topic_id in agent.opinions else 0
            for agent in env.agents_active.values()
        ) / n_agents
        logging.info(f"Average opinion on '{topic_name}' of all agents: {avg_opinion:.2f}")
    # List groups and their members
    for group in env.groups_active.values():
        logging.info(f"\n{group}")
        for member in group.list_members():
            logging.info(f"  - {member}")
    for iteration in range(n_iterations):
        logging.info(f"\n--- Communication round {iteration+1} ---")    
        # Each agent in the negative group communicates with a random neutral agent
        for agent in negative.members:
            other_agent = random.choice(list(neutral.members))
            agent.communicate(other_agent.id)
        # Each agent in the positive group communicates with a random neutral agent
        for agent in positive.members:
            other_agent = random.choice(list(neutral.members))
            agent.communicate(other_agent.id)
        # Record opinions after this round
        opinions_over_time.append([agent.opinions.copy() for agent in env.agents_active.values()])
        # Log the average opinion of all agents in the environment after communication
        avg_opinions = {}
        for topic_name, topic_id in topic_name_to_id.items():
            avg_opinion = sum(
                agent.opinions[topic_id].value if topic_id in agent.opinions else 0
                for agent in env.agents_active.values()
            ) / n_agents
            avg_opinions[topic_name] = avg_opinion
            logging.info(f"Average opinion on '{topic_name}' of all agents after communication: {avg_opinion:.2f}")
    logging.info("\nAgent communication demo complete.")
    # --- Plotting ---
    output_dir = Path("data/output")
    output_dir.mkdir(parents=True, exist_ok=True)
    colors = {"negative": "lightcoral", "neutral": "lightblue", "positive": "lightgreen"}
    for topic_name, topic_id in topic_name_to_id.items():
        plt.figure(figsize=(8, 5))
        topic_opinions = [
            [agent[topic_id].value if topic_id in agent else 0 for agent in opinions]
            for opinions in opinions_over_time
        ]
        plt.boxplot(topic_opinions, positions=range(len(opinions_over_time)), patch_artist=True, boxprops=dict(facecolor=colors[topic_name]), medianprops=dict(color='red'))
        plt.xlabel('Round')
        plt.ylabel(f"Opinions ({topic_name})")
        plt.title(f"Distribution of '{topic_name.capitalize()}' Opinions Over Time")
        plt.xticks(range(len(opinions_over_time)), [f"{i}" for i in range(len(opinions_over_time))])
        plt.grid(axis='y', linestyle='--', alpha=0.7)
        plt.tight_layout()
        filename = output_dir / f"opinions_{topic_name}.png"
        plt.savefig(filename)
        plt.close()
        logging.info(f"Boxplot of '{topic_name}' opinions saved to {filename}")

def run_survey():
    """
    Run a survey to change agent opinions.
    """
    logging.info("Survey mode selected...")

    # Initialise
    year = 2026
    env = Environment(year=year)
    raise NotImplementedError("Survey mode is not yet implemented.")

def main():
    parser = argparse.ArgumentParser(description="Run GABM simulation.")
    parser.add_argument('--mode', choices=['example', 'survey'], default='example',
                        help="Select run mode: 'example' for the demo, 'survey' to change agent opinions.")
    args = parser.parse_args()

    if args.mode == 'example':
        run_example()
    elif args.mode == 'survey':
        run_survey()

if __name__ == "__main__":
    # Set up logging to file and console
    log_dir = Path("data/logs")
    log_dir.mkdir(parents=True, exist_ok=True)
    log_file = log_dir / "run_main.log"
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        handlers=[
            logging.FileHandler(log_file, mode="w"),
            logging.StreamHandler(sys.stdout)
        ]
    )
    main()
