BASE_PATH = "/Volumes/workspace/default/contextual_mesh"

df_sales = spark.read.option("header", True).option("inferSchema", True).csv(
    f"{BASE_PATH}/silver_product_sales_monthly.csv"
)
df_context = spark.read.option("header", True).option("inferSchema", True).csv(
    f"{BASE_PATH}/silver_context_monthly.csv"
)
df_rules = spark.read.option("header", True).option("inferSchema", True).csv(
    f"{BASE_PATH}/gold_expert_rules.csv"
)
df_ground_truth = spark.read.option("header", True).option("inferSchema", True).csv(
    f"{BASE_PATH}/gold_ground_truth_validation.csv"
)
df_sources = spark.read.option("header", True).option("inferSchema", True).csv(
    f"{BASE_PATH}/source_catalog.csv"
)

evidence = (
    df_rules.alias("r")
    .join(
        df_context.alias("c"),
        F.col("r.ContextFactor") == F.col("c.ContextFactor"),
        "inner"
    )
    .join(
        df_sales.alias("s"),
        (
            (F.col("r.Product") == F.col("s.Product")) &
            (F.col("r.ProductLine") == F.col("s.ProductLine")) &
            (F.col("c.Month") == F.col("s.Month")) &
            (F.col("c.Market") == F.col("s.Market"))
        ),
        "inner"
    )
    .select(
        F.col("c.Month").alias("Month"),
        F.col("c.Market").alias("Market"),
        F.col("r.ProductLine").alias("ProductLine"),
        F.col("r.Product").alias("Product"),
        F.col("r.ContextFactor").alias("ContextFactor"),
        F.col("c.CurrentValue").alias("ContextValue"),
        F.col("s.Units").alias("Units"),
        F.col("s.Revenue").alias("Revenue"),
        F.col("s.GrossProfit").alias("GrossProfit")
    )
)

pearson_index = (
    evidence
    .groupBy("ProductLine", "Product", "ContextFactor")
    .agg(
        F.count("*").alias("Observations"),
        F.corr("ContextValue", "Units").alias("Pearson_Units"),
        F.corr("ContextValue", "Revenue").alias("Pearson_Revenue"),
        F.corr("ContextValue", "GrossProfit").alias("Pearson_GrossProfit")
    )
)

window_spec = (
    Window
    .partitionBy("ProductLine", "Product", "ContextFactor")
    .orderBy("Month")
)

for lag in range(0, 7):
    lagged_evidence = lagged_evidence.withColumn(
        f"ContextValue_Lag_{lag}",
        F.lag("ContextValue", lag).over(window_spec)
        if lag > 0 else F.col("ContextValue")
    )

# Compute one correlation per rule and lag, then choose
# the largest absolute correlation.

regularized_lag_index = (
    lag_index
    .withColumn("AbsCorrelation", F.abs("Correlation"))
    .withColumn(
        "LagPenalty",
        F.lit(1.0) - F.col("LagMonths") * F.lit(0.05)
    )
    .withColumn(
        "RegularizedScore",
        F.col("AbsCorrelation") * F.col("LagPenalty")
    )
)

context_wide = (
    df_context
    .groupBy("Month", "Market")
    .pivot("ContextFactor")
    .agg(F.first("CurrentValue"))
)

model_base = df_sales.join(
    context_wide,
    on=["Month", "Market"],
    how="inner"
)

X_scaled = StandardScaler().fit_transform(X)
y_scaled = StandardScaler().fit_transform(
    y.to_numpy().reshape(-1, 1)
).ravel()

model = LinearRegression()
model.fit(X_scaled, y_scaled)

pipeline = Pipeline([
    ("scaler", StandardScaler()),
    (
        "model",
        ElasticNetCV(
            l1_ratio=[0.1, 0.25, 0.5, 0.75, 0.9, 1.0],
            alphas=np.logspace(-4, 1, 100),
            cv=5,
            max_iter=100000,
            random_state=42
        )
    )
])

pipeline.fit(X, y)

consensus_index = (
    method_outputs
    .withColumn(
        "PositiveVotes",
        sum(
            F.when(F.col(c) == "POSITIVE", 1).otherwise(0)
            for c in direction_columns
        )
    )
    .withColumn(
        "NegativeVotes",
        sum(
            F.when(F.col(c) == "NEGATIVE", 1).otherwise(0)
            for c in direction_columns
        )
    )
    .withColumn(
        "ConsensusDirection",
        F.when(F.col("PositiveVotes") > F.col("NegativeVotes"), "POSITIVE")
         .when(F.col("NegativeVotes") > F.col("PositiveVotes"), "NEGATIVE")
         .otherwise("CONFLICT")
    )
)