from typing import Union
import pathlib 
from pathlib import Path
import pandas as pd
import matplotlib.pyplot as plt  
import numpy as np
import math


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.file_name = self._verify_filename(data_fname) # Verify the file name
        self.data_fname= pathlib.Path(data_fname) # The path of the wanted file
   
    def _verify_filename(self, data_fname):
        """ Verify the input filename. """
        file_path=pathlib.Path(data_fname)

        if (type(data_fname) != str) and (type(data_fname) != pathlib.WindowsPath): # Check if the file name is a string or a path
            raise TypeError ("Wrong input")   
        elif file_path.exists()==False: # Check if the file name exists
            raise ValueError ("Filename does not exists.")

        return data_fname

    def read_data(self):
        """Reads the json data located in self.data_fname into memory, to
        the attribute self.data.
                """
        data_df = pd.read_json(self.data_fname) # Load the file
        self.data=data_df # Saving the data to a the attribute self.data
        
        return self.data

    def show_age_distrib(self):
        """Calculates and plots the age distribution of the participants.

        Returns
        -------
        hist : np.ndarray
        Number of people in a given bin
        bins : np.ndarray
        Bin edges
        """
        data_df=self.read_data() # Put the data into a data frame
        hist_values=data_df['age'] # Take the values of the age
        (n, bins, patches) =plt.hist(hist_values, bins=[0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100])

        return n, bins

    def remove_rows_without_mail(self) -> pd.DataFrame:
        """Checks self.data for rows with invalid emails, and removes them.

        Returns
        -------
        df : pd.DataFrame
        A corrected DataFrame, i.e. the same table but with the erroneous rows removed and
        the (ordinal) index after a reset.
        """

        data_df=self.read_data() # Put the data into a data frame
        emails=data_df['email'] # Creat a series of the email column 
        check_list=[] # Creat a list for email validity output

        for mail in emails: # Check the mails
            check_list.append(('@' in mail) and (('@.' in mail)==False) and (('.' in mail) and (mail.find('.')<len(mail))))
            
        corrected_df = data_df.loc[check_list] # Creat a new Data frame
        corrected_with_r_ind=corrected_df.reset_index(drop=True)

        return corrected_with_r_ind
    
    def fill_na_with_mean(self):
        """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
        -------
        df : pd.DataFrame
        The corrected DataFrame after insertion of the mean grade
        arr : np.ndarray
        Row indices of the students that their new grades were generated
        """

        data_df=self.read_data() # Put the data into a data frame
        questions_ans=data_df.loc[:, 'q1':'q5'] # Indexes of the wanted columns 

        rows_with_nan=pd.isnull(questions_ans).any(axis=1) # Finds rows with NaNs
        
        nan_indexes=questions_ans.loc[rows_with_nan].index
        indexes_array=np.array(nan_indexes) # Array of the subjects with NaNs

        nan_subject_data=questions_ans.loc[rows_with_nan]

        new_df=pd.DataFrame(index =indexes_array, columns =['q1', 'q2','q3','q4','q5'])

        for ind in indexes_array: # Replace the NaN's by a mean and put it into a new data frame
            new_row=nan_subject_data.loc[ind].fillna(nan_subject_data.loc[ind].mean(axis=0))
            new_df.loc[ind]=new_row
        
        return new_df, indexes_array


    def score_subjects(self, maximal_nans_per_sub: int = 1) -> pd.DataFrame:
        """Calculates the average score of a subject and adds a new "score" column
        with it.

        If the subject has more than "maximal_nans_per_sub" NaN in his grades, the
        score should be NA. Otherwise, the score is simply the mean of the other grades.
        The datatype of score is UInt8, and the floating point raw numbers should be
        rounded down.

        Parameters
        ----------
        maximal_nans_per_sub : int, optional
        Number of allowed NaNs per subject before giving a NA score.

        Returns
        -------
        pd.DataFrame
        A new DF with a new column - "score".
        """

        data_df=self.read_data() # Put the data into a data frame
        questions_ans=data_df.loc[:, 'q1':'q5'] # Indexes of the wanted columns
        number_of_nan=questions_ans.isnull().sum(axis=1).tolist()
        floored_mean_score=np.floor(questions_ans.mean(axis=1))#.astype(pd.UInt8Dtype()))

        for idx, num in enumerate(number_of_nan):
             if num>maximal_nans_per_sub:
                floored_mean_score[idx]=None
             else:
                 continue
        
        data_df['score']=floored_mean_score.astype(pd.UInt8Dtype())

        return data_df

    
    def correlate_gender_age(self) -> pd.DataFrame:
        """Looks for a correlation between the gender of the subject, their age
        and the score for all five questions.
        Returns
        -------
        pd.DataFrame
        A DataFrame with a MultiIndex containing the gender and whether the subject is above
        40 years of age, and the average score in each of the five questions.
        """

        data_df=self.read_data()
        data_df.set_index(['gender', 'age'],append=True) # Creats multi index based on rows gender and age
        data_df.groupby(['gender', 'age']).groups

        return data_df



StarcoderdataPython