from .data_process import cleanData
import altair as alt
import pandas as pd
import ibis


class graphGenerator(cleanData):
    def _init_(
        self, saving_dir: str = "data/", database_url: str = "duckdb:///data.ddb"
    ):
        super().__init__(saving_dir, database_url)

    def create_graph(self, naics_code: str) -> alt.Chart:
        naics_data = self.group_by_naics_code()

        filtered_df = naics_data.filter(naics_data["first_4_naics_code"] == naics_code)

        # Joining the Year and the quarter columns for graphing purpose
        filtered_df = filtered_df.mutate(
            year_qtr=(
                filtered_df["year"].cast("string")
                + ibis.literal("-Q")
                + filtered_df["qtr"].cast("string")
            )
        )
        filtered_pd = filtered_df.execute()

        filtered_pd = filtered_pd.sort_values(by=["year", "qtr"])

        chart = (
            alt.Chart(filtered_pd)
            .mark_line()
            .encode(
                x=alt.X("year_qtr", title="Year", sort=list(filtered_pd["year_qtr"])),
                y=alt.Y("total_employment_sum:Q", title="Total Employment"),
                tooltip=[
                    alt.Tooltip("year_qtr", title="Year and Quarter"),
                    alt.Tooltip("total_employment_sum", title="Total Employment Sum"),
                ],
            )
            .properties(
                title="Employment Trends for NAICS 5412", width=1000, height=400
            )
        )

        return chart

    def gen_naics_graph(self, naics_code: str) -> alt.Chart:
        df_filtered, naics = self.get_naics_data(naics_code)

        line = (
            alt.Chart(df_filtered)
            .mark_line()
            .encode(
                x=alt.X("x_axis:N", title="Year"),
                y=alt.Y("total_employment_sum:Q", title="Total Employment"),
                tooltip=["x_axis", "total_employment_sum"],
            )
        )
        points = (
            alt.Chart(df_filtered)
            .mark_point(color="darkblue", size=60, filled=True)
            .encode(
                x="x_axis:N",
                y="total_employment_sum:Q",
                tooltip=["x_axis", "total_employment_sum"],
            )
        )
        chart = (
            (line + points)
            .properties(
                title=f"Employment in the US by Quarter for NAICS {naics_code}",
                width=1000,
                height=200,
            )
            .configure_view(fill="#e6f7ff")
            .configure_axis(gridColor="white", grid=True)
            .configure_title(anchor="start", fontSize=16, color="#333333", offset=30)
        )
        context = {
            "naics_code": naics,
        }
        return chart, context

    def gen_wages_graph(
        self, time_frame: str, naics_desc: str, data_type: str, selected_column: str
    ) -> alt.Chart:
        if data_type == "nivel":
            column = selected_column
        elif data_type == "primera_diferencia":
            if selected_column == "average_salary":
                column = "salary_diff"
            elif selected_column == "total_wages":
                column = "payroll_diff"
            else:
                column = f"{selected_column}_diff"
        elif data_type == "cambio_porcentual":
            if selected_column == "average_salary":
                column = "salary_diff_p"
            elif selected_column == "total_wages":
                column = "payroll_diff_p"
            else:
                column = f"{selected_column}_diff_p"
        df, naics = self.filter_wages_data(time_frame, naics_desc, column)

        columns = [
            "taxable_wages",
            "total_wages",
            "average_salary",
            "social_security",
            "medicare",
            "contributions_due",
        ]
        columns = [
            {"value": col, "label": col.replace("_", " ").capitalize()}
            for col in columns
        ]

        x_values = df.select("time_period").unique().to_series().to_list()

        if time_frame == "quarterly":
            tick_vals = x_values[::3]
        else:
            tick_vals = x_values

        chart = (
            alt.Chart(df)
            .mark_line()
            .encode(
                x=alt.X("time_period:N", title="", axis=alt.Axis(values=tick_vals)),
                y=alt.Y("nominas:Q", title=""),
                tooltip=[
                    alt.Tooltip("time_period", title="Time Period"),
                    alt.Tooltip("nominas", title="Wages"),
                ],
            )
            .properties(
                width="container",
            )
            .configure_view(fill="#e6f7ff")
            .configure_axis(gridColor="white", grid=True)
        )

        return chart, naics, columns


if __name__ == "__main__":
    g = graphGenerator()
    print(g.create_graph())
