import csv
import warnings
import pandas as pd
from pandas import DataFrame

import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.backends.backend_qt5agg import FigureCanvas 

from PyQt6 import uic, QtWidgets
from PyQt6.QtGui import QWindow
from PyQt6.QtWidgets import QDialog, QApplication, QTableWidget, \
                            QTableWidgetItem, QHeaderView

class TitanicPassengers(QWindow):
    """
    The main application window.
    """
    
    def __init__(self):
        """
        Load the UI and initialize its components.
        """
        super().__init__()
        warnings.simplefilter(action='ignore', category=FutureWarning)
                
        self._passenger_df = pd.read_csv('TitanicSurvival.csv')
        self.ui = uic.loadUi('titanic_passengers.ui')
        
        self._initialize_class_menu()
        self._initialize_passenger_table()
        self.ui.count_label.setText('')
        
        # Connect event handlers to the menu and the check boxes.
        self.ui.class_menu.currentIndexChanged.connect(
                                        self._enter_passenger_data)
        self.ui.checkbox_1st.stateChanged.connect(self._draw_bar_chart)
        self.ui.checkbox_2nd.stateChanged.connect(self._draw_bar_chart)
        self.ui.checkbox_3rd.stateChanged.connect(self._draw_bar_chart)
        
        self.ui.show();
        
    def _initialize_class_menu(self):
        """
        Initialize the passenger class menu.
        """
        self.ui.class_menu.addItem('[select passenger class]', '')
        self.ui.class_menu.addItem('All passengers', 'all')
        self.ui.class_menu.addItem('1st class passengers', '1st')
        self.ui.class_menu.addItem('2nd class passengers', '2nd')
        self.ui.class_menu.addItem('3rd class passengers', '3rd')
        
    def _initialize_passenger_table(self):
        """
        Initialize the passenger table.
        """
        self.ui.passenger_table.clear()
        
        columns = ['Survived', 'Sex', 'Age', 'Name']
        self.ui.passenger_table.setHorizontalHeaderLabels(columns)
        
    def _enter_passenger_data(self):    
        """
        Enter passenger data from the dataframe into 
        the passenger table.
        """    
        klass = self.ui.class_menu.currentData()
        
        # Determine which passenger class.
        if klass == '':
            self._initialize_passenger_table()
            self.ui.count_label.setText('')
            return
        elif klass == 'all':
            df = self._passenger_df
        else:
            df = self._passenger_df[
                self._passenger_df.passengerClass == klass]
               
        row_count = len(df.index)
                
        # Set the passenger data into the table cells.
        self.ui.passenger_table.setRowCount(row_count)        
        for row_index in range(len(df.index)):
            row = df.iloc[row_index]
            col_index = 0
            
            for data in [row.survived, row.sex, row.age, row[0]]:
                item = QTableWidgetItem(str(data))
                self.ui.passenger_table.setItem(row_index, 
                                                col_index, 
                                                item)
                col_index += 1
        
        # Display the row count.
        self.ui.count_label.setText(f'{row_count:,d}')
        
    def _draw_bar_chart(self):
        """
        Draw a Seaborn bar chart of passenger survival
        and set it in the graph layout widget.
        """
        self._clear_graph()
            
        # Read the check boxes to determine for which 
        # passenger classes to draw bars.
        pclass = []
        if self.ui.checkbox_1st.isChecked():
            pclass.append('1st')
        if self.ui.checkbox_2nd.isChecked():
            pclass.append('2nd')
        if self.ui.checkbox_3rd.isChecked():
            pclass.append('3rd')
        
        if len(pclass) == 0:
            return

        # Create a table for the chosen passenger classes.
        table  = []        
        for yn in ['yes', 'no']:
            for klass in pclass:
                count = len(self._passenger_df[self._passenger_df \
                                .passengerClass == klass] \
                                .survived[self._passenger_df \
                                .survived == yn])
                table.append([klass, count, yn])
        
        df = DataFrame(table)
        df.columns = ['Class', 'Count', 'Survived']

        # Graph the dataframe and set the graph 
        # in the graph layout widget.
        sns.set(style='whitegrid')
        ax = sns.barplot(x='Class', y='Count', hue='Survived',
                         data=df,  errorbar=None, 
                         palette=['lawngreen', 'saddlebrown'])
        self.ui.graph_layout.addWidget(FigureCanvas(ax.figure))

        plt.close()
        
    def _clear_graph(self):
        """
        Remove all the bars from the graph.
        """
        children = []
        
        # Gather children which are the bars in the layout.
        for i in range(self.ui.graph_layout.count()):
            child = self.ui.graph_layout.itemAt(i).widget()
            if child:
                children.append(child)
                
        # Delete the bars.
        for child in children:
            child.deleteLater()