import pandas as pd
import os
import sys
import random

class Automat:
    def __init__(self, df: pd.DataFrame, end_states: list):
        self._df = df
        self._end_states = end_states
    
    def _consume_character(self, state: int, char: str):
        filter = (self._df[0] == state) & (self._df[2] == char)
        matching_rows = self._df[filter]
        if matching_rows.empty:
            return None
        end_states = matching_rows.iloc[:,1].tolist()
        return random.choice(end_states)
    
    def consume_word(self, word: str):
        prev_state = 0
        state = 0
        for char in word:
            prev_state = state
            state = self._consume_character(state, char)
            if state is None:
                return prev_state in self._end_states
        return state in self._end_states

def create_automat_from_file(file_path: str):
    automat_df = pd.read_csv(file_path, sep='\t', header=None)
    
    end_states_condition = (automat_df[1].isnull()) & (automat_df[2].isnull())
    end_states_df = automat_df[end_states_condition]
    end_states = end_states_df.iloc[:,0].tolist()
    
    automat_df = automat_df[~end_states_condition]
    automat_df[2] = automat_df[2].astype(str)

    return Automat(automat_df, end_states)

automat_path = os.path.join(os.path.dirname(__file__), 'test1.arg')
automat = create_automat_from_file(automat_path)

test_path = os.path.join(os.path.dirname(__file__), 'test1.in')


results = []
with open(test_path, 'r') as f:
    for line in f:
        results.append('YES' if automat.consume_word(line.strip()) else 'NO')

out_path = os.path.join(os.path.dirname(__file__), 'test1.out')

with open(out_path, 'w') as f:
    for result in results:
        f.write(result + '\n')