# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# # Pyspark for Linear Regression

# ### 1. Set up spark context and SparkSession

# +
from pyspark.sql import SparkSession

spark = SparkSession \
    .builder \
    .appName("Python Spark SQL basic example") \
    .config("spark.some.config.option", "some-value") \
    .getOrCreate()
# -

# ### 2. Load dataset

df = spark.read.format('com.databricks.spark.csv').\
                               options(header='true', \
                               inferschema='true').load("./data/Advertising.csv",header=True);

df.take(2)
df.printSchema()

df.show(6)

# ### 3. Convert the data to features (dense vector) and label

from pyspark.sql import Row
from pyspark.ml.linalg import Vectors


# convert the data to dense vector
#def transData(row):
#    return Row(label=row["Sales"],
#               features=Vectors.dense([row["TV"],
#                                       row["Radio"],
#                                       row["Newspaper"]]))
def transData(data):
    return data.rdd.map(lambda r: [Vectors.dense(r[:-1]),r[-1]]).toDF(['features','label'])


# convert the data to dense vector
#def transData(row):
#    return Row(label=row["Sales"],
#               features=Vectors.dense([row["TV"],
#                                       row["Radio"],
#                                       row["Newspaper"]]))
def transData(data):
    return data.rdd.map(lambda r: [Vectors.dense(r[:-1]),r[-1]]).toDF(['features','label'])


# ### 4. Transform the dataset to DataFrame

#transformed = df.rdd.map(transData).toDF() 
data= transData(df)
data.show(6)

# ### 5. Convert features data format and set up training and test data sets

from pyspark.ml import Pipeline
from pyspark.ml.regression import LinearRegression
from pyspark.ml.feature import VectorIndexer
from pyspark.ml.evaluation import RegressionEvaluator

featureIndexer = VectorIndexer(inputCol="features", \
                               outputCol="indexedFeatures",\
                               maxCategories=4).fit(data)

# Split the data into training and test sets (40% held out for testing)
(trainingData, testData) = data.randomSplit([0.6, 0.4])

# ### 5. Fit model (Ridge Regression and the LASSO)

lr = LinearRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8)

# Chain indexer and tree in a Pipeline
pipeline = Pipeline(stages=[featureIndexer, lr])

# Train model.  This also runs the indexer.
model = pipeline.fit(trainingData)

 lrmodel= model.stages[1]

lrmodel.coefficients

# or 
[stage.coefficients for stage in model.stages if hasattr(stage, "coefficients")]

lrmodel.summary.meanAbsoluteError

# ### 6. Make predictions

predictions = model.transform(testData)

# Select example rows to display.
predictions.select("prediction", "label", "features").show(5)

# ### 8. Evaluation

# Select (prediction, true label) and compute test error
evaluator = RegressionEvaluator(
    labelCol="label", predictionCol="prediction", metricName="rmse")
rmse = evaluator.evaluate(predictions)
print("Root Mean Squared Error (RMSE) on test data = %g" % rmse)

# - summary

temp_path = 'temp/Users/wenqiangfeng/Dropbox/Spark/Code/model'
modelPath = temp_path + "/lr_model"
model.save(modelPath)

# - save and extract model

lr2 = model.load(modelPath)

lr2.coefficients

StarcoderdataJupyterScriptsDedupFiltered