from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import StructType, StructField, StringType, DoubleType
from pyspark.ml.linalg import Vectors, VectorUDT
from pyspark.ml.classification import DecisionTreeClassifier
from pyspark.ml.feature import StringIndexer, VectorAssembler
import pandas as pd
from statsmodels.stats.outliers_influence import variance_inflation_factor
from .utils import replace_feature_indices_with_names
from .utils import parse_tree_structure
from .utils import round_off

import graphviz # type: ignore

@udf(returnType=VectorUDT())
def to_dense_vector(v):
    return Vectors.dense(v.toArray())

def get_multivariate_analysis(df, table_name, numerical_columns,id_columns=None, vif_analysis=1, decision_tree_analysis=1, target_column=None, depth=3):
    spark = SparkSession.builder.getOrCreate()

    # Drop the ID column if provided
    if id_columns:
        df = df.drop(*id_columns)

    vif_schema = StructType([
        StructField('Feature', StringType(), True),
        StructField('VIF', DoubleType(), True)
    ])
    # Drop the table if it exists
    spark.sql(f"DROP TABLE IF EXISTS {table_name}")

    if vif_analysis and numerical_columns:
        for column in numerical_columns:
            try:
                # Convert Spark DataFrame to Pandas DataFrame
                df_pd = df.select(numerical_columns).dropna().toPandas()

                # Calculate VIF for the current column
                vif = variance_inflation_factor(df_pd.values, df_pd.columns.get_loc(column))

                new_row = spark.createDataFrame([(column, round_off(vif))], schema=vif_schema)
                new_row.write.option("mergeSchema", "true").saveAsTable(table_name, mode='append')
            except Exception as e:
                print(f"Error processing VIF for column {column}: {e}")
        print(f"The results have been successfully saved to the table: {table_name}")

    if decision_tree_analysis and numerical_columns and target_column:
        try:
            string_indexer = StringIndexer(inputCol=target_column, outputCol="indexed_target")
            si_model = string_indexer.fit(df)
            df_indexed = si_model.transform(df)

            assembler_inputs = [col for col in numerical_columns if col != target_column]
            feature_assembler = VectorAssembler(inputCols=assembler_inputs, outputCol="features")
            df_features = feature_assembler.transform(df_indexed)

            df_features = df_features.withColumn("features_dense", to_dense_vector("features"))

            dt = DecisionTreeClassifier(labelCol="indexed_target", featuresCol="features_dense", maxDepth=depth)
            dt_model = dt.fit(df_features)

            print("Decision Tree Model of", target_column, ":")
            tree_model_readable = replace_feature_indices_with_names(dt_model.toDebugString, assembler_inputs)
            print(tree_model_readable)

            # Generate Graphviz representation
            dot_string = parse_tree_structure(tree_model_readable)

            # Visualize the decision tree using Graphviz
            graph = graphviz.Source(dot_string)
            graph.render("decision_tree", format="png")
            print("Decision tree rendered to decision_tree.png")
        except Exception as e:
            print(f"Error processing Decision Tree analysis: {e}")