import java.util.*;
import java.io.*;

public class PoliticalSurvey {
  private static final String[] POLITICAL_PARTIES = {
      "Democratic", "Republican", "Libertarian", "Green"
  };

  private static final Map<Integer, Map<Integer, Map<String, Double>>> partyResponsePatterns = new HashMap<>();

  // let's create a training data points per party
  private static final int[] partyDataPoints = new int[POLITICAL_PARTIES.length];

  // I gathered some random politiacl questions
  private static final List<SurveyQuestion> QUESTIONS = new ArrayList<>();

  private static final Scanner scanner = new Scanner(System.in);

  private static final double CONFIDENCE_THRESHOLD = 0.65;

  public static void main(String[] args) {
    initializeQuestions();
    initializePartyResponsePatterns();

    System.out.println("==========================================================");
    System.out.println("    POLITICAL PARTY PREDICTION SURVEY");
    System.out.println("==========================================================");
    System.out.println("This program will try to predict your political affiliation");
    System.out.println("based on your responses to several questions.");
    System.out.println();

    boolean continueSurvey = true;

    while (continueSurvey) {
      List<String> userResponses = conductSurvey();
      addTrainingData(userResponses);

      System.out.println("\nWould you like to take the survey again? (yes/no)");
      String response = scanner.nextLine().trim().toLowerCase();
      continueSurvey = response.equals("yes") || response.equals("y");

      if (continueSurvey) {
        System.out.println("\n==========================================================\n");
      }
    }

    System.out.println("\nThank you for using the Political Party Prediction Survey!");
    scanner.close();
  }

  // Initialize the survey questions
  private static void initializeQuestions() {
    QUESTIONS.add(new SurveyQuestion(
        "What should the government do to help the poor?",
        new String[] {
            "Make it easier to apply for assistance",
            "Allow parents to use education funds for charter schools",
            "Create welfare-to-work programs",
            "Nothing, private charity should handle it"
        }));

    QUESTIONS.add(new SurveyQuestion(
        "What is your view on taxation?",
        new String[] {
            "Higher taxes on the wealthy to fund social programs",
            "Lower taxes to stimulate economic growth",
            "Minimal taxation, only for essential services",
            "Tax policies focused on environmental sustainability"
        }));

    QUESTIONS.add(new SurveyQuestion(
        "What is your stance on healthcare?",
        new String[] {
            "Universal healthcare funded by the government",
            "Market-based healthcare with limited government involvement",
            "Free-market healthcare with no government involvement",
            "Community-based healthcare systems with government support"
        }));

    QUESTIONS.add(new SurveyQuestion(
        "What is your view on gun control?",
        new String[] {
            "Stricter gun control laws are needed",
            "Existing gun laws are sufficient",
            "Gun ownership is a fundamental right with minimal restrictions",
            "Gun control with community-based oversight"
        }));

    QUESTIONS.add(new SurveyQuestion(
        "How should we address climate change?",
        new String[] {
            "Government regulations and investment in renewable energy",
            "Market-based solutions with limited regulations",
            "No government intervention, let the market decide",
            "Aggressive environmental policies and global cooperation"
        }));

    QUESTIONS.add(new SurveyQuestion(
        "What should be our approach to foreign policy?",
        new String[] {
            "Global cooperation and diplomatic engagement",
            "Strong national defense and protecting American interests",
            "Non-interventionism and free trade",
            "International cooperation for peace and environmental protection"
        }));

    QUESTIONS.add(new SurveyQuestion(
        "What is your stance on immigration?",
        new String[] {
            "Path to citizenship for undocumented immigrants",
            "Stronger border security and merit-based immigration",
            "Open borders with minimal restrictions",
            "Humanitarian-focused immigration policies"
        }));

    QUESTIONS.add(new SurveyQuestion(
        "Which political party do you most closely affiliate with?",
        POLITICAL_PARTIES));
  }

  // Initialize response patterns with starting weights, just please keep in mind
  // that these weights are not 100% accurate.
  private static void initializePartyResponsePatterns() {
    // Democratic Party patterns
    Map<Integer, Map<String, Double>> democraticPatterns = new HashMap<>();
    democraticPatterns.put(0, createResponseWeights(new double[] { 0.7, 0.1, 0.2, 0.0 }));
    democraticPatterns.put(1, createResponseWeights(new double[] { 0.7, 0.1, 0.0, 0.2 }));
    democraticPatterns.put(2, createResponseWeights(new double[] { 0.8, 0.1, 0.0, 0.1 }));
    democraticPatterns.put(3, createResponseWeights(new double[] { 0.8, 0.1, 0.0, 0.1 }));
    democraticPatterns.put(4, createResponseWeights(new double[] { 0.7, 0.1, 0.0, 0.2 }));
    democraticPatterns.put(5, createResponseWeights(new double[] { 0.7, 0.2, 0.0, 0.1 }));
    democraticPatterns.put(6, createResponseWeights(new double[] { 0.7, 0.1, 0.0, 0.2 }));
    partyResponsePatterns.put(0, democraticPatterns);

    // Republican Party patterns
    Map<Integer, Map<String, Double>> republicanPatterns = new HashMap<>();
    republicanPatterns.put(0, createResponseWeights(new double[] { 0.1, 0.4, 0.4, 0.1 }));
    republicanPatterns.put(1, createResponseWeights(new double[] { 0.0, 0.8, 0.2, 0.0 }));
    republicanPatterns.put(2, createResponseWeights(new double[] { 0.0, 0.8, 0.2, 0.0 }));
    republicanPatterns.put(3, createResponseWeights(new double[] { 0.1, 0.5, 0.4, 0.0 }));
    republicanPatterns.put(4, createResponseWeights(new double[] { 0.0, 0.7, 0.3, 0.0 }));
    republicanPatterns.put(5, createResponseWeights(new double[] { 0.0, 0.8, 0.2, 0.0 }));
    republicanPatterns.put(6, createResponseWeights(new double[] { 0.0, 0.8, 0.1, 0.1 }));
    partyResponsePatterns.put(1, republicanPatterns);

    // Libertarian Party patterns
    Map<Integer, Map<String, Double>> libertarianPatterns = new HashMap<>();
    libertarianPatterns.put(0, createResponseWeights(new double[] { 0.0, 0.2, 0.1, 0.7 }));
    libertarianPatterns.put(1, createResponseWeights(new double[] { 0.0, 0.3, 0.7, 0.0 }));
    libertarianPatterns.put(2, createResponseWeights(new double[] { 0.0, 0.2, 0.8, 0.0 }));
    libertarianPatterns.put(3, createResponseWeights(new double[] { 0.0, 0.1, 0.9, 0.0 }));
    libertarianPatterns.put(4, createResponseWeights(new double[] { 0.0, 0.2, 0.8, 0.0 }));
    libertarianPatterns.put(5, createResponseWeights(new double[] { 0.0, 0.1, 0.9, 0.0 }));
    libertarianPatterns.put(6, createResponseWeights(new double[] { 0.0, 0.1, 0.9, 0.0 }));
    partyResponsePatterns.put(2, libertarianPatterns);

    // Green Party patterns
    Map<Integer, Map<String, Double>> greenPatterns = new HashMap<>();
    greenPatterns.put(0, createResponseWeights(new double[] { 0.4, 0.0, 0.1, 0.5 }));
    greenPatterns.put(1, createResponseWeights(new double[] { 0.3, 0.0, 0.0, 0.7 }));
    greenPatterns.put(2, createResponseWeights(new double[] { 0.6, 0.0, 0.0, 0.4 }));
    greenPatterns.put(3, createResponseWeights(new double[] { 0.4, 0.0, 0.0, 0.6 }));
    greenPatterns.put(4, createResponseWeights(new double[] { 0.3, 0.0, 0.0, 0.7 }));
    greenPatterns.put(5, createResponseWeights(new double[] { 0.3, 0.0, 0.0, 0.7 }));
    greenPatterns.put(6, createResponseWeights(new double[] { 0.2, 0.0, 0.0, 0.8 }));
    partyResponsePatterns.put(3, greenPatterns);

    // Initialize partyDataPoints with fictional data points
    // to simulate pre-existing training data
    partyDataPoints[0] = 100; // Democratic
    partyDataPoints[1] = 100; // Republican
    partyDataPoints[2] = 50; // Libertarian
    partyDataPoints[3] = 50; // Green
  }

  // Helper method to create response weights map
  private static Map<String, Double> createResponseWeights(double[] weights) {
    Map<String, Double> responseWeights = new HashMap<>();
    for (int i = 0; i < weights.length; i++) {
      responseWeights.put(String.valueOf(i), weights[i]);
    }
    return responseWeights;
  }

  // Add this new method to save survey responses to a TXT file
  private static void saveSurveyToTxt(List<String> userResponses, String predictedParty, double confidence,
      String actualParty) {
    try {
      // Create a filename with timestamp
      String filename = "survey_response_" + ".txt";

      try (PrintWriter writer = new PrintWriter(new FileWriter(filename))) {
        writer.println("POLITICAL SURVEY RESPONSE");
        writer.println("=========================");
        writer.println();

        for (int i = 0; i < QUESTIONS.size() - 1; i++) {
          SurveyQuestion question = QUESTIONS.get(i);
          String answer = question.getOptions()[Integer.parseInt(userResponses.get(i))];

          writer.println("Question " + (i + 1) + ": " + question.getQuestion());
          writer.println("Answer: " + answer);
          writer.println();
        }

        writer.println("PREDICTION RESULTS");
        writer.println("-----------------");
        writer.println("Predicted Party: " + predictedParty + " Party");
        writer.println("Confidence: " + String.format("%.1f%%", confidence * 100));
        writer.println("Actual Party: " + actualParty + " Party");
        writer.println();

        writer.println("Survey completed successfully!");
      }

      System.out.println("\nSurvey responses have been saved to: " + filename);

    } catch (IOException e) {
      System.err.println("Error saving survey responses to file: " + e.getMessage());
    }
  }

  private static List<String> conductSurvey() {
    List<String> userResponses = new ArrayList<>();
    int questionNum = 0;

    int lastQuestionIndex = QUESTIONS.size() - 1;

    String predictedParty = "";
    double highestConfidence = 0.0;

    System.out.println("Please answer the following questions:");

    while (questionNum < lastQuestionIndex) {
      SurveyQuestion question = QUESTIONS.get(questionNum);

      System.out.println("\nQuestion " + (questionNum + 1) + ": " + question.getQuestion());

      String[] options = question.getOptions();
      for (int i = 0; i < options.length; i++) {
        System.out.println(((char) ('A' + i)) + ". " + options[i]);
      }

      System.out.print("\nYour answer (A-" + (char) ('A' + options.length - 1) + "): ");
      String userInput = scanner.nextLine().trim().toUpperCase();

      if (userInput.length() == 1 && userInput.charAt(0) >= 'A' && userInput.charAt(0) < 'A' + options.length) {
        int optionIndex = userInput.charAt(0) - 'A';
        userResponses.add(String.valueOf(optionIndex));

        // After collecting some responses, try to predict party
        if (questionNum >= 2) {
          double[] partyConfidences = calculatePartyConfidences(userResponses);

          // Find the party with highest confidence
          int maxConfidencePartyIndex = 0;
          highestConfidence = partyConfidences[0];

          for (int i = 1; i < partyConfidences.length; i++) {
            if (partyConfidences[i] > highestConfidence) {
              maxConfidencePartyIndex = i;
              highestConfidence = partyConfidences[i];
            }
          }

          predictedParty = POLITICAL_PARTIES[maxConfidencePartyIndex];

          if (highestConfidence >= CONFIDENCE_THRESHOLD) {
            System.out.println("\nBased on your answers so far, I predict you might align with the " +
                predictedParty + " Party (Confidence: " +
                String.format("%.1f%%", highestConfidence * 100) + ")");
          }
        }

        questionNum++;
      } else {
        System.out.println("Invalid input. Please enter a letter between A and " +
            (char) ('A' + options.length - 1));
      }
    }

    if (!predictedParty.isEmpty()) {
      System.out.println("\n==========================================================");
      System.out.println("FINAL PREDICTION: Based on your responses, I predict your");
      System.out.println("political affiliation is " + predictedParty + " Party");
      System.out.println("Confidence: " + String.format("%.1f%%", highestConfidence * 100));
      System.out.println("==========================================================\n");
    }

    SurveyQuestion lastQuestion = QUESTIONS.get(lastQuestionIndex);
    System.out.println("Final Question: " + lastQuestion.getQuestion());

    String[] partyOptions = lastQuestion.getOptions();
    for (int i = 0; i < partyOptions.length; i++) {
      System.out.println(((char) ('A' + i)) + ". " + partyOptions[i]);
    }

    boolean validInput = false;
    while (!validInput) {
      System.out.print("\nYour answer (A-" + (char) ('A' + partyOptions.length - 1) + "): ");
      String userInput = scanner.nextLine().trim().toUpperCase();

      if (userInput.length() == 1 && userInput.charAt(0) >= 'A' && userInput.charAt(0) < 'A' + partyOptions.length) {
        int optionIndex = userInput.charAt(0) - 'A';
        userResponses.add(String.valueOf(optionIndex));
        validInput = true;

        // We ompare prediction with actual affiliation
        String actualParty = partyOptions[optionIndex];
        System.out.println("\nYour actual affiliation: " + actualParty + " Party");

        if (!predictedParty.isEmpty()) {
          if (actualParty.equals(predictedParty)) {
            System.out.println("Prediction was CORRECT!");
          } else {
            System.out.println("Prediction was incorrect. This data will help improve future predictions.");
          }
        }

        saveSurveyToTxt(userResponses, predictedParty, highestConfidence, actualParty);
      } else {
        System.out.println("Invalid input. Please enter a letter between A and " +
            (char) ('A' + partyOptions.length - 1));
      }
    }

    return userResponses;
  }

  // here we calculate confidence scores for each party based on user responses
  private static double[] calculatePartyConfidences(List<String> userResponses) {
    double[] partyScores = new double[POLITICAL_PARTIES.length];

    for (int partyIndex = 0; partyIndex < POLITICAL_PARTIES.length; partyIndex++) {
      Map<Integer, Map<String, Double>> partyPattern = partyResponsePatterns.get(partyIndex);
      double totalScore = 0.0;

      for (int questionIndex = 0; questionIndex < userResponses.size(); questionIndex++) {
        String response = userResponses.get(questionIndex);
        Map<String, Double> questionPatterns = partyPattern.get(questionIndex);

        if (questionPatterns != null && questionPatterns.containsKey(response)) {
          double weight = questionPatterns.get(response);
          totalScore += weight;
        }
      }

      partyScores[partyIndex] = totalScore / userResponses.size();
    }

    double sum = 0.0;
    for (double score : partyScores) {
      sum += score;
    }

    if (sum > 0) {
      for (int i = 0; i < partyScores.length; i++) {
        partyScores[i] /= sum;
      }
    }

    return partyScores;
  }

  // Add new training data based on user responses
  private static void addTrainingData(List<String> userResponses) {
    if (userResponses.size() != QUESTIONS.size()) {
      return;
    }

    int actualPartyIndex = Integer.parseInt(userResponses.get(userResponses.size() - 1));
    partyDataPoints[actualPartyIndex]++;

    Map<Integer, Map<String, Double>> partyPattern = partyResponsePatterns.get(actualPartyIndex);

    for (int questionIndex = 0; questionIndex < userResponses.size() - 1; questionIndex++) {
      String response = userResponses.get(questionIndex);

      Map<String, Double> questionPattern = partyPattern.computeIfAbsent(questionIndex, k -> new HashMap<>());

      for (int optionIndex = 0; optionIndex < QUESTIONS.get(questionIndex).getOptions().length; optionIndex++) {
        String option = String.valueOf(optionIndex);
        double currentWeight = questionPattern.getOrDefault(option, 0.0);

        if (option.equals(response)) {
          double newWeight = (currentWeight * (partyDataPoints[actualPartyIndex] - 1) + 1.0) /
              partyDataPoints[actualPartyIndex];
          questionPattern.put(option, newWeight);
        } else {
          if (currentWeight > 0) {
            double newWeight = (currentWeight * (partyDataPoints[actualPartyIndex] - 1)) /
                partyDataPoints[actualPartyIndex];
            questionPattern.put(option, newWeight);
          }
        }
      }

      double sum = 0.0;
      for (double weight : questionPattern.values()) {
        sum += weight;
      }

      if (sum > 0) {
        for (String option : questionPattern.keySet()) {
          questionPattern.put(option, questionPattern.get(option) / sum);
        }
      }
    }
  }

  private static class SurveyQuestion {
    private String question;
    private String[] options;

    public SurveyQuestion(String question, String[] options) {
      this.question = question;
      this.options = options;
    }

    public String getQuestion() {
      return question;
    }

    public String[] getOptions() {
      return options;
    }
  }
}