from typing import Union,Tuple
import pathlib
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("TkAgg")
from matplotlib import pyplot as plt

class QuestionnaireAnalysis:
    """
    Reads and analyzes data generated by the questionnaire experiment.
    Should be able to accept strings and pathlib.Path objects.
    """
    
    def __init__(self, data_fname: Union[pathlib.Path, str]):
        self.data_fname = data_fname
        self.data = self.read_data()
        
    def read_data(self):
        """
        Reads the json data located in self.data_fname into memory, to
        the attribute self.data.
        """
        with open(self.data_fname) as f:
            data = pd.read_json(f)
        return data        

    def show_age_distrib(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculates and plots the age distribution of the participants.
        Returns a tuple containing two numpy arrays:
        The first item being the number of people in a given bin.
        The second item being the bin edges.
        """
        data = self.data['age'].as_matrix()
        (n,bins,patches) = plt.hist(data, bins= range(0,100,10))
        plt.show()
        return (n,bins)

    def remove_rows_without_mail(self) -> pd.DataFrame:
        """
        Checks self.data for rows with invalid emails, and removes them.
        Returns the corrected DataFrame, i.e. the same table but with
        the erroneous rows removed and the (ordinal) index after a reset.
        """
        data = self.data
        to_drop = []
        for item, frame in data['email'].iteritems():
            print(0,item,1,frame)
            if frame.endswith('.com') == False or '@' in frame == False:
                to_drop.append(item)
        data_emails = data.drop(index=to_drop)
        data_emails = data_emails.reset_index(drop=True)
        return data_emails 

    def fill_na_with_mean(self) -> Union[pd.DataFrame, np.ndarray]:
        """
        Finds, in the original DataFrame, the subjects that didn't answer
        all questions, and replaces that missing value with the mean of the
        other grades for that student. Returns the corrected DataFrame,
        as well as the row indices of the students that their new grades
        werdatae generated.
        """
        data = self.data
        filled_data = data.fillna(data[['q1','q2','q3','q4','q5']].mean())
        rows = np.where(data[['q1','q2','q3','q4','q5']].isnull())
        return filled_data, rows[0]
if __name__ == "__main__":
    test = QuestionnaireAnalysis('data.json')
    test.show_age_distrib()
    test.remove_rows_without_mail()
    test.fill_na_with_mean()
