main.cpp

#include <iostream>
#include <string>
#include <vector>
#include "DecisionTree.h"
#include "LinearRegression.h"
#include "StockData.h"
using namespace std;

//Global variables for testing purposes
const int maxIterations = 1000;
const double learningRate = 0.01;


int main() {
    string FolderInput;
    string FileInput;

    // First prompt: ETF or Stock selection
    cout << "Would you like to load an ETF(1), or Stock(2): ";
    cin >> FolderInput;

    // Second prompt: Specific file selection
    cout << "Which specific entry would you like to load (which .txt file): ";
    cin >> FileInput;

    string filePath;
    if (FolderInput == "1") {
        filePath = "../ETFs/" + FileInput;
    } else if (FolderInput == "2") {
        filePath = "../Stocks/" + FileInput;
    } else {
        cout << "Invalid input" << endl;
        return 1;
    }

    dataLoader dataLoaderObj;
    vector<stockDataPoint> allData = dataLoaderObj.readFile(filePath);
    pair<vector<stockDataPoint>, vector<stockDataPoint>> splitResult = dataLoaderObj.splitData(allData);
    vector<stockDataPoint> trainingData = splitResult.first;  // 80% training data
    vector<stockDataPoint> testData = splitResult.second;  // 20% test data
    dataLoaderObj.scaleData(trainingData, testData);
    cout << "Data loaded successfully" << endl;

    LinearRegression linearRegression;
    DecisionTree decisionTree;
    string menuInput;
    while (true) {
        cout << "===============================" << endl;
        cout << "\033[31mStock Price Prediction System\033[0m" << endl;
        cout << "===============================" << endl;

        cout << "1. Train Linear Regression" << endl;
        cout << "2. Train Decision Tree" << endl;
        cout << "3. Compare Performance" << endl;
        cout << "4. Predict using Linear Regression" << endl;
        cout << "5. Predict using Decision Tree" << endl;
        cout << "6. Exit" << endl;
        cout << "Enter your choice: ";
        cin >> menuInput;

        if (menuInput == "6") {
            cout << "Exiting program" << endl;
            return 0;
        }
        else if (menuInput == "1") {
            cout << "Training Linear Regression" << endl;
            linearRegression.trainLinearRegression(trainingData, learningRate, maxIterations);
            cout << "Linear Regression trained successfully" << endl;
        }
        else if (menuInput == "2") {
            cout << "Training Decision Tree" << endl;
            auto [trainFeatures, trainTargets] = dataLoaderObj.extractFeaturesAndTargets(trainingData);
            decisionTree.train(trainFeatures, trainTargets);
            cout << "Decision Tree trained successfully" << endl;
        }
        else if (menuInput == "3") {
            cout << "Comparing Performance" << endl;

            // Calculate Linear Regression performance using testData
            vector<double> predictions;
            vector<double> actual;

            for (const auto& testPoint : testData) {
                float prediction = linearRegression.predict(testPoint.stockInputCategories);
                predictions.push_back((double)prediction);
                actual.push_back((double)testPoint.target);
            }

            double linearRegressionMSE = linearRegression.MSE(predictions, actual);
            cout << "Linear Regression Performance (MSE): " << linearRegressionMSE << endl;

            // Test Decision Tree on test data
            auto [testFeatures, testTargets] = dataLoaderObj.extractFeaturesAndTargets(testData);
            vector<double> decisionTreePredictions = decisionTree.predictBatch(testFeatures);

            // Calculate MSE for Decision Tree
            double decisionTreeMSE = linearRegression.MSE(decisionTreePredictions, testTargets);
            cout << "Decision Tree Performance (MSE): " << decisionTreeMSE << endl;
            ofstream mseFile("mse_results.csv");
            if (mseFile.is_open()) {
                mseFile << "Model,MSE\n";
                mseFile << "LinearRegression," << linearRegressionMSE << "\n";
                mseFile << "DecisionTree," << decisionTreeMSE << "\n";
                mseFile.close();
                cout << "MSE results exported to mse_results.csv" << endl;
            } 
            else{
                cerr << "Failed to write MSE results file." << endl;
            }
            return 0;
        }
        else if (menuInput == "4") {
            // Predict using Linear Regression
            double open, high, low;
            int volume, openInt;

            cout << "Enter open price: ";
            cin >> open;
            cout << "Enter high price: ";
            cin >> high;
            cout << "Enter low price: ";
            cin >> low;
            cout << "Enter volume: ";
            cin >> volume;
            cout << "Enter openInt: ";
            cin >> openInt;

            // Construct feature vector with bias term first (same structure as training data)
            vector<double> rawInput;
            //features.push_back(1.0);  // Bias term
            rawInput.push_back(open);
            rawInput.push_back(high);
            rawInput.push_back(low);
            rawInput.push_back((double)volume);
            rawInput.push_back((double)openInt);

            vector<double> scaledInput = dataLoaderObj.scaleSingleInput(rawInput);
            vector<double> features;
            features.push_back(1.0);
            features.insert(features.end(), scaledInput.begin(), scaledInput.end());

            double prediction = linearRegression.predict(features);
            cout << "Predicted close price: " << prediction << endl;
        }
        else if (menuInput == "5") {
            // Predict using Decision Tree
            double open, high, low;
            int volume, openInt;

            cout << "Enter open price: ";
            cin >> open;
            cout << "Enter high price: ";
            cin >> high;
            cout << "Enter low price: ";
            cin >> low;
            cout << "Enter volume: ";
            cin >> volume;
            cout << "Enter openInt: ";
            cin >> openInt;

            // Construct feature vector with bias term first (same structure as training data)
            vector<double> features;
            features.push_back(1.0);  // Bias term
            features.push_back(open);
            features.push_back(high);
            features.push_back(low);
            features.push_back((double)volume);
            features.push_back((double)openInt);

            double prediction = decisionTree.predict(features);
            cout << "Predicted close price: " << prediction << endl;
        }

        else {
            cout << "Invalid input" << endl;
            return 0;
        }
    }

}

// TIP See CLion help at <a
// href="https://www.jetbrains.com/help/clion/">jetbrains.com/help/clion/</a>.
//  Also, you can try interactive lessons for CLion by selecting
//  'Help | Learn IDE Features' from the main menu.