import json
import requests
import random
import string
import time
import re
import collections
from urllib.parse import parse_qs, urlencode, urlparse

# Machine Learning Libraries
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

# Disable SSL warnings (optional, based on your environment)
from requests.packages.urllib3.exceptions import InsecureRequestWarning
requests.packages.urllib3.disable_warnings(InsecureRequestWarning)

# Define vowels for potential prioritization
vowels = ['a', 'e', 'i', 'o', 'u']

# Load and clean the dictionary
def load_dictionary(file_path):
    try:
        with open(file_path, "r") as f:
            words = [line.strip().lower() for line in f if line.strip().isalpha()]
        return words
    except FileNotFoundError:
        print(f"Dictionary file not found at {file_path}. Please verify the path.")
        return []

class HangmanAPIError(Exception):
    """Custom exception for Hangman API errors."""
    def __init__(self, result):
        self.result = result
        self.code = None
        try:
            self.type = result["error_code"]
        except (KeyError, TypeError):
            self.type = ""
        try:
            self.message = result["error_description"]
        except (KeyError, TypeError):
            try:
                self.message = result["error"]["message"]
                self.code = result["error"].get("code")
                if not self.type:
                    self.type = result["error"].get("type", "")
            except (KeyError, TypeError):
                try:
                    self.message = result["error_msg"]
                except (KeyError, TypeError):
                    self.message = "Unknown Hangman API Error."
        super().__init__(self.message)

class HangmanAPI(object):
    def __init__(self, access_token=None, session=None, timeout=10):
        self.hangman_url = self.determine_hangman_url()
        self.access_token = access_token
        self.session = session or requests.Session()
        self.timeout = timeout
        self.guessed_letters = set()  # Use set to prevent duplicates
        
        # Update the path to your dictionary file as needed
        full_dictionary_location = "/kaggle/input/hangmen/words_250000_train.txt"  # Ensure this path is correct
        self.full_dictionary = self.build_dictionary(full_dictionary_location)
        if not self.full_dictionary:
            raise ValueError("Failed to load the dictionary. Terminating the program.")
        print(f"Loaded dictionary with {len(self.full_dictionary)} words.")
        
        # Precompute letter frequencies in the full dictionary
        self.full_dictionary_common_letter_sorted = collections.Counter("".join(self.full_dictionary)).most_common()
        
        # Initialize current dictionary
        self.current_dictionary = self.full_dictionary.copy()
        
        # Statistics
        self.total_practice_runs = 0
        self.total_recorded_runs = 0
        self.total_recorded_successes = 0
        self.total_practice_successes = 0

    @staticmethod
    def determine_hangman_url():
        """Determine the best Hangman API URL based on response times."""
        links = ['https://trexsim.com', 'https://sg.trexsim.com']
        data = {link: 0 for link in links}
        
        print("Determining the fastest Hangman API URL...")
        for link in links:
            try:
                s = time.time()
                response = requests.get(link, timeout=5)
                if response.status_code == 200:
                    elapsed = time.time() - s
                    data[link] = elapsed
                    print(f"Link {link} responded in {elapsed:.4f} seconds.")
                else:
                    data[link] = float('inf')
                    print(f"Link {link} responded with status code {response.status_code}.")
            except requests.RequestException as e:
                data[link] = float('inf')
                print(f"Link {link} failed with exception: {e}")
                
        sorted_links = sorted(data.items(), key=lambda x: x[1])
        best_link = sorted_links[0][0] if sorted_links and sorted_links[0][1] != float('inf') else 'https://trexsim.com'
        hangman_endpoint = '/trexsim/hangman'
        final_url = best_link + hangman_endpoint
        print(f"Determined Hangman URL: {final_url}")
        return final_url

    def build_dictionary(self, file_path):
        """Load the dictionary from the specified file."""
        words = load_dictionary(file_path)
        if words:
            print(f"Dictionary loaded from {file_path}.")
        return words

    def request(self, path, args=None, post_args=None, method=None):
        """Send a request to the Hangman API."""
        if args is None:
            args = {}
        if post_args is not None:
            method = "POST"

        # Add access_token to the request
        if self.access_token:
            if post_args and "access_token" not in post_args:
                post_args["access_token"] = self.access_token
            elif "access_token" not in args:
                args["access_token"] = self.access_token

        time.sleep(0.2)  # Rate limiting

        num_retry, time_sleep = 50, 2
        for attempt in range(num_retry):
            try:
                response = self.session.request(
                    method or "GET",
                    self.hangman_url + path,
                    timeout=self.timeout,
                    params=args,
                    json=post_args,
                    verify=False
                )
                response.raise_for_status()
                break
            except requests.HTTPError as e:
                try:
                    error_content = e.response.text
                    error_json = json.loads(error_content)
                    raise HangmanAPIError(error_json)
                except json.JSONDecodeError:
                    print(f"Non-JSON error response received: {e.response.text}")
                    raise HangmanAPIError({"error_msg": "HTTP Error with no content."})
            except requests.exceptions.SSLError as e:
                if attempt + 1 == num_retry:
                    raise
                print(f"SSL Error encountered. Retrying in {time_sleep} seconds...")
                time.sleep(time_sleep)
            except requests.RequestException as e:
                if attempt + 1 == num_retry:
                    raise
                print(f"Request exception encountered: {e}. Retrying in {time_sleep} seconds...")
                time.sleep(time_sleep)

        headers = response.headers
        content_type = headers.get('content-type', '')
        if 'application/json' in content_type:
            try:
                result = response.json()
            except json.JSONDecodeError:
                print(f"Failed to decode JSON response. Response Text: {response.text}")
                raise HangmanAPIError({"error_msg": "Failed to decode JSON response."})
        else:
            try:
                result = json.loads(response.text)
            except json.JSONDecodeError:
                print(f"Non-JSON response received: {response.text}")
                raise HangmanAPIError({"error_msg": "Response content type not supported."})
        
        if isinstance(result, dict) and result.get("error"):
            print(f"API Error: {result.get('error')}")
            raise HangmanAPIError(result)
        return result

    def build_state(self, word, guessed_letters, tries_left):
        """Create a hashable state from the current word, guessed letters, and tries left."""
        clean_word = word.replace(" ", "").lower()  # Remove spaces and lowercase
        state = f"{clean_word}|{''.join(sorted(guessed_letters))}|{tries_left}"
        return state

    def start_game(self, player, practice=True, verbose=True):
        """
        Start a new game of Hangman using the provided player (RLAgent).
        """
        # Reset guessed letters and dictionary
        self.guessed_letters = set()
        self.current_dictionary = self.full_dictionary.copy()
        
        # Initialize status
        status = "ongoing"  # Set initial status
        
        # Start new game with GET method and 'practice' parameter
        try:
            response = self.request("/new_game", args={"practice": int(practice)}, post_args=None, method="GET")
        except HangmanAPIError as e:
            if verbose:
                print(f"HangmanAPIError during new_game: {e.message}")
            return False
        except Exception as e:
            if verbose:
                print(f"Unexpected error during new_game: {e}")
            return False

        if response.get('status') == "approved":
            game_id = response.get('game_id')
            word = response.get('word', "")  # Example: "_ p p _ e "
            tries_remains = response.get('tries_remains')
            
            if verbose:
                print(f"New game started. Game ID: {game_id}. Tries left: {tries_remains}. Word: {word}")
            
            # Update statistics
            if practice:
                self.total_practice_runs += 1
            else:
                self.total_recorded_runs += 1
            
            state = self.build_state(word, self.guessed_letters, tries_remains)
            
            while tries_remains > 0 and status != "success":
                # Agent selects an action (guess letter)
                guess_letter = player.select_action(word)
                
                if not guess_letter:
                    if verbose:
                        print("No available letters to guess.")
                    break
                if guess_letter in self.guessed_letters:
                    # Penalize for repeating a guess
                    player.store_rewards(-0.1)
                    if verbose:
                        print(f"Repeated guess: '{guess_letter}'. Penalizing.")
                    continue
                
                self.guessed_letters.add(guess_letter)
                
                if verbose:
                    print(f"Guessing letter: '{guess_letter}'")
                
                # Make the guess and get response using GET method
                try:
                    res = self.request(
                        "/guess_letter",
                        args={
                            "request": "guess_letter",
                            "game_id": game_id,
                            "letter": guess_letter
                        },
                        post_args=None,
                        method="GET"
                    )
                except HangmanAPIError as e:
                    if verbose:
                        print(f"HangmanAPIError during guess_letter: {e.message}")
                    break
                except Exception as e:
                    if verbose:
                        print(f"Unexpected error during guess_letter: {e}")
                    break
                
                if verbose:
                    print(f"Server response: {res}")
                
                status = res.get('status', 'failed')  # Default to 'failed' if not provided
                tries_remains = res.get('tries_remains', tries_remains)
                correct_guess = res.get('correct_guess', False)
                new_word = res.get('word', word).lower()
                
                if correct_guess:
                    # Reward for correct guess
                    player.store_rewards(0.1)
                    if verbose:
                        print(f"Correct guess: '{guess_letter}'. Reward +0.1")
                else:
                    # Penalty for incorrect guess
                    player.store_rewards(-0.05)
                    if verbose:
                        print(f"Incorrect guess: '{guess_letter}'. Penalty -0.05")
                
                if status == "success":
                    # Reward for winning the game
                    player.store_rewards(1.0)
                    if verbose:
                        print("Game won!")
                    if practice:
                        self.total_practice_successes += 1
                    else:
                        self.total_recorded_successes += 1
                    break
                elif status == "failed":
                    # Penalty for losing the game
                    player.store_rewards(-1.0)
                    if verbose:
                        print("Game lost.")
                    break
                elif status == "ongoing":
                    word = new_word

            # Update policy after each game
            player.update_policy()
            return status == "success"
        else:
            if verbose:
                print("Failed to start a new game.")
                print(f"Response: {response}")
            return False

    class PolicyNetwork(nn.Module):
        def __init__(self, input_size, hidden_size, output_size, num_layers=2):
            super(PolicyNetwork, self).__init__()
            self.hidden_size = hidden_size
            self.num_layers = num_layers

            # Embedding layer for input characters
            self.embedding = nn.Embedding(input_size, hidden_size)

            # LSTM layer
            self.lstm = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True)

            # Fully connected layer to output probabilities for each letter
            self.fc = nn.Linear(hidden_size, output_size)

        def forward(self, x, hidden):
            # x: (batch_size, seq_length)
            embedded = self.embedding(x)  # (batch_size, seq_length, hidden_size)
            lstm_out, hidden = self.lstm(embedded, hidden)  # lstm_out: (batch_size, seq_length, hidden_size)
            lstm_out = lstm_out[:, -1, :]  # Take the output of the last time step
            output = self.fc(lstm_out)  # (batch_size, output_size)
            return output, hidden

        def init_hidden(self, batch_size=1):
            # Initialize hidden state and cell state to zeros
            weight = next(self.parameters()).data
            hidden = (weight.new_full((self.num_layers, batch_size, self.hidden_size), 0.0),
                      weight.new_full((self.num_layers, batch_size, self.hidden_size), 0.0))
            return hidden

    class RLAgent:
        def __init__(self, alphabet, lr=1e-3, hidden_size=128, num_layers=2, gamma=0.99, entropy_weight=0.01):
            self.alphabet = alphabet
            self.char_to_ix = {ch: idx for idx, ch in enumerate(self.alphabet)}
            self.ix_to_char = {idx: ch for idx, ch in enumerate(self.alphabet)}
            self.input_size = len(alphabet)
            self.output_size = len(alphabet)
            self.hidden_size = hidden_size
            self.num_layers = num_layers
            self.gamma = gamma
            self.entropy_weight = entropy_weight

            self.policy_net = PolicyNetwork(self.input_size, self.hidden_size, self.output_size, self.num_layers)
            self.optimizer = optim.Adam(self.policy_net.parameters(), lr=lr)
            self.reset()

        def reset(self):
            self.saved_log_probs = []
            self.saved_entropies = []
            self.rewards = []

        def select_action(self, state):
            """
            state: string representing the current word state (e.g., "_ p p _ e")
            """
            # Convert state to a sequence of indices
            state_seq = self.state_to_indices(state)
            state_tensor = torch.tensor(state_seq, dtype=torch.long).unsqueeze(0)  # (1, seq_length)

            hidden = self.policy_net.init_hidden(batch_size=1)
            output, hidden = self.policy_net(state_tensor, hidden)
            probs = torch.softmax(output, dim=1)
            m = torch.distributions.Categorical(probs)
            action = m.sample()
            self.saved_log_probs.append(m.log_prob(action))
            self.saved_entropies.append(m.entropy())
            guessed_letter = self.ix_to_char[action.item()]
            return guessed_letter

        def state_to_indices(self, state):
            """
            Convert the current state string to a list of character indices.
            Replace unknown letters with a special character (e.g., '_').
            """
            chars = list(state.replace(" ", ""))
            indices = []
            for ch in chars:
                if ch in self.char_to_ix:
                    indices.append(self.char_to_ix[ch])
                else:
                    indices.append(self.char_to_ix['_'])  # Use '_' index for unknown letters
            return indices

        def update_policy(self):
            """
            Update policy network using the collected rewards and log probabilities.
            Incorporates entropy regularization to encourage exploration.
            """
            R = 0
            policy_losses = []
            returns = []

            # Discount rewards
            for r in self.rewards[::-1]:
                R = r + self.gamma * R
                returns.insert(0, R)
            
            returns = torch.tensor(returns)
            # Normalize returns to reduce variance
            if returns.std() != 0:
                returns = (returns - returns.mean()) / (returns.std() + 1e-5)
            else:
                returns = returns - returns.mean()
            
            for log_prob, entropy, R in zip(self.saved_log_probs, self.saved_entropies, returns):
                # Policy gradient objective
                policy_losses.append(-log_prob * R)
                # Entropy regularization to encourage exploration
                policy_losses.append(-self.entropy_weight * entropy)
            
            self.optimizer.zero_grad()
            policy_loss = torch.cat(policy_losses).sum()
            policy_loss.backward()
            self.optimizer.step()

            self.reset()

        def store_rewards(self, reward):
            self.rewards.append(reward)

    class RLHangmanPlayer:
        def __init__(self, api, agent, verbose=False):
            self.api = api
            self.agent = agent
            self.verbose = verbose

        def play_game(self, practice=True):
            # Reset guessed letters and dictionary
            self.api.guessed_letters = set()
            self.api.current_dictionary = self.api.full_dictionary.copy()
            
            # Initialize status
            status = "ongoing"  # Set initial status

            # Start new game with GET method and 'practice' parameter
            try:
                response = self.api.request("/new_game", args={"practice": int(practice)}, post_args=None, method="GET")
            except HangmanAPIError as e:
                if self.verbose:
                    print(f"HangmanAPIError during new_game: {e.message}")
                return False
            except Exception as e:
                if self.verbose:
                    print(f"Unexpected error during new_game: {e}")
                return False

            if response.get('status') == "approved":
                game_id = response.get('game_id')
                word = response.get('word', "")  # Example: "_ p p _ e "
                tries_remains = response.get('tries_remains')
                
                if self.verbose:
                    print(f"New game started. Game ID: {game_id}. Tries left: {tries_remains}. Word: {word}")
                
                # Update statistics
                if practice:
                    self.api.total_practice_runs += 1
                else:
                    self.api.total_recorded_runs += 1
                
                state = self.api.build_state(word, self.api.guessed_letters, tries_remains)
                self.agent.reset()

                while tries_remains > 0 and status != "success":
                    # Agent selects an action (guess letter)
                    guess_letter = self.agent.select_action(word)
                    
                    if not guess_letter:
                        if self.verbose:
                            print("No available letters to guess.")
                        break
                    if guess_letter in self.api.guessed_letters:
                        # Penalize for repeating a guess
                        self.agent.store_rewards(-0.1)
                        if self.verbose:
                            print(f"Repeated guess: '{guess_letter}'. Penalizing.")
                        continue

                    self.api.guessed_letters.add(guess_letter)
                    
                    if self.verbose:
                        print(f"Guessing letter: '{guess_letter}'")
                    
                    # Make the guess and get response using GET method
                    try:
                        res = self.api.request(
                            "/guess_letter",
                            args={
                                "request": "guess_letter",
                                "game_id": game_id,
                                "letter": guess_letter
                            },
                            post_args=None,
                            method="GET"
                        )
                    except HangmanAPIError as e:
                        if self.verbose:
                            print(f"HangmanAPIError during guess_letter: {e.message}")
                        break
                    except Exception as e:
                        if self.verbose:
                            print(f"Unexpected error during guess_letter: {e}")
                        break
                    
                    if self.verbose:
                        print(f"Server response: {res}")
                    
                    status = res.get('status', 'failed')  # Default to 'failed' if not provided
                    tries_remains = res.get('tries_remains', tries_remains)
                    correct_guess = res.get('correct_guess', False)
                    new_word = res.get('word', word).lower()
                    
                    if correct_guess:
                        # Reward for correct guess
                        self.agent.store_rewards(0.1)
                        if self.verbose:
                            print(f"Correct guess: '{guess_letter}'. Reward +0.1")
                    else:
                        # Penalty for incorrect guess
                        self.agent.store_rewards(-0.05)
                        if self.verbose:
                            print(f"Incorrect guess: '{guess_letter}'. Penalty -0.05")
                    
                    if status == "success":
                        # Reward for winning the game
                        self.agent.store_rewards(1.0)
                        if self.verbose:
                            print("Game won!")
                        if practice:
                            self.api.total_practice_successes += 1
                        else:
                            self.api.total_recorded_successes += 1
                        break
                    elif status == "failed":
                        # Penalty for losing the game
                        self.agent.store_rewards(-1.0)
                        if self.verbose:
                            print("Game lost.")
                        break
                    elif status == "ongoing":
                        word = new_word

                # Update policy after each game
                self.agent.update_policy()
                return status == "success"
            else:
                if self.verbose:
                    print("Failed to start a new game.")
                    print(f"Response: {response}")
                return False

    # Main Execution Block
    if __name__ == "__main__":
        # Initialize HangmanAPI with your access token
        api = HangmanAPI(access_token="775976eaf5c7653d7f9cc3f826ad3e", timeout=10)
        
        # Initialize RL Agent
        # Alphabet includes all lowercase letters plus '_' for unknown letters
        alphabet = string.ascii_lowercase + '_'
        agent = RLAgent(alphabet=alphabet, lr=1e-3, hidden_size=128, num_layers=2, gamma=0.99, entropy_weight=0.01)
        
        # Initialize RL Hangman Player
        player = RLHangmanPlayer(api=api, agent=agent, verbose=False)
        
        # Training parameters
        num_training_games = 1000  # Adjust based on available resources
        print_every = 100  # How often to print status updates

        successes = 0

        for game in range(1, num_training_games + 1):
            success = player.play_game(practice=True)
            if success:
                successes += 1
            
            if game % print_every == 0:
                current_success_rate = successes / game
                print(f'--- After {game} Games ---')
                print(f'Total practice runs: {api.total_practice_runs}')
                print(f'Practice success rate: {api.total_practice_successes / api.total_practice_runs:.2%} if api.total_practice_runs > 0 else 0.0%')
                print(f'Current success rate (last {print_every}): {current_success_rate:.2%}\n')
            
            # Ensure not exceeding 20 new games per minute
            time.sleep(3)  # Approximately 20 games per minute (60 / 20 = 3 seconds per game)
        
        # After training, you can play additional games or evaluate the agent further
        print("Training completed.")
        
        # Optionally, save the model
        torch.save(agent.policy_net.state_dict(), "policy_network.pth")
        
        # To start recorded games (once finalized), uncomment and use cautiously
        """
        for i in range(1, 1001):
            print(f'Playing {i}th recorded game')
            player.play_game(practice=False)
            time.sleep(3)  # Adjust sleep time according to rate limits
        """
        
        # Retrieve and display final status after all games
        status = api.my_status()
        if isinstance(status, list) and len(status) == 4:
            total_practice_runs, total_recorded_runs, total_recorded_successes, total_practice_successes = status
        else:
            # Fallback to internal counters if the response is not as expected
            total_practice_runs = api.total_practice_runs
            total_recorded_runs = api.total_recorded_runs
            total_recorded_successes = api.total_recorded_successes
            total_practice_successes = api.total_practice_successes
        
        # Calculate final success rates safely
        if total_recorded_runs > 0:
            recorded_success_rate = total_recorded_successes / total_recorded_runs
        else:
            recorded_success_rate = 0.0
        
        if total_practice_runs > 0:
            practice_success_rate = total_practice_successes / total_practice_runs
        else:
            practice_success_rate = 0.0
        
        print(f'--- Final Statistics ---')
        print(f'Total practice runs: {total_practice_runs}')
        print(f'Practice success rate: {practice_success_rate:.2%}')
        print(f'Overall success rate = {recorded_success_rate:.2%}')
this code is not working properly check it and optimise it