const candlestickData = require('./data.js');

class TradingEnvironment {
    constructor(data) {
        this.data = data;
        console.log(`Total number of candles: ${this.data.length}`);
        this.initialCash = 10000;
        this.reset();
    }

    reset() {
        this.currentStep = 0;
        this.cash = this.initialCash;
        this.sharesOwned = 0;
        this.done = false;
        return this.getState();
    }

    getState() {
        if (this.currentStep >= this.data.length) {
            this.done = true;
            return null;
        }

        const candle = this.data[this.currentStep];
        
        return [
            parseFloat(candle[1]), // open
            parseFloat(candle[2]), // high
            parseFloat(candle[3]), // low
            parseFloat(candle[4]), // close
            parseFloat(candle[5])  // volume
        ];
    }

    step(action) {
        if (this.done) {
            console.log("Episode already finished");
            return { state: null, reward: 0, done: true };
        }

        const currentPrice = parseFloat(this.data[this.currentStep][4]); // Closing price

        // Apply action
        if (action === 1) { // Buy
            
            const maxShares = Math.floor(this.cash / currentPrice);
            this.sharesOwned += maxShares;
            this.cash -= maxShares * currentPrice;
            console.log('Bought! Cash is ' + this.cash)
        } else if (action === 2) { // Sell
            this.cash += this.sharesOwned * currentPrice;
            this.sharesOwned = 0;
            console.log('Sold! Cash is ' + this.cash)
        }

        // Calculate reward before moving to next step
        const portfolioValue = this.cash + (this.sharesOwned * currentPrice);
        const reward = portfolioValue - this.initialCash;

        // Move to next step
        this.currentStep++;
        if (this.currentStep >= this.data.length) {
            this.done = true;
        }

        const nextState = this.getState();

        return { state: nextState, reward, done: this.done };
    }

    getCurrentPortfolioValue() {
        if (this.currentStep >= this.data.length) {
            return this.cash; // If we're at the end, just return the cash
        }
        const currentPrice = parseFloat(this.data[this.currentStep][4]);
        return this.cash + (this.sharesOwned * currentPrice);
    }
}


  const tf = require('@tensorflow/tfjs');

// Define the Q-network (policy network)
function createQNetwork(inputShape, outputSize) {
  const model = tf.sequential();
  model.add(tf.layers.dense({ units: 128, inputShape: inputShape, activation: 'relu' }));
  model.add(tf.layers.dense({ units: 128, activation: 'relu' }));
  model.add(tf.layers.dense({ units: outputSize }));  // Output Q-values for each action (Buy, Sell, Hold)
  model.compile({ optimizer: 'adam', loss: 'meanSquaredError' });
  return model;
}

// Experience Replay (to improve learning)
class ReplayBuffer {
  constructor(size) {
    this.buffer = [];
    this.size = size;
  }

  add(experience) {
    if (this.buffer.length >= this.size) {
      this.buffer.shift();  // Remove the oldest experience
    }
    this.buffer.push(experience);
  }

  sample(batchSize) {
    const sample = [];
    for (let i = 0; i < batchSize; i++) {
      const index = Math.floor(Math.random() * this.buffer.length);
      sample.push(this.buffer[index]);
    }
    return sample;
  }
}

// Q-Learning Agent
class DQNAgent {
  constructor(stateSize, actionSize) {
    this.stateSize = stateSize;
    this.actionSize = actionSize;
    this.model = createQNetwork([stateSize], actionSize);
    this.replayBuffer = new ReplayBuffer(2000);
    this.gamma = 0.95;  // Discount factor for future rewards
    this.epsilon = 1.0;  // Exploration rate (start by exploring)
    this.epsilonDecay = 0.995;
    this.epsilonMin = 0.01;
  }

  // Select an action based on an epsilon-greedy policy
  selectAction(state) {
    if (Math.random() < this.epsilon) {
      return Math.floor(Math.random() * this.actionSize);  // Random action
    } else {
      const qValues = this.model.predict(tf.tensor2d([state])).dataSync();
      return qValues.indexOf(Math.max(...qValues));  // Action with the highest Q-value
    }
  }

  // Train the model with a batch of experiences
  train(batch) {
    for (const [state, action, reward, nextState, done] of batch) {
        let target = reward;
        if (!done) {
            const futureQ = this.model.predict(tf.tensor2d([nextState])).max().dataSync()[0];
            target += this.gamma * futureQ;
        }

        const qValues = this.model.predict(tf.tensor2d([state]));
        qValues.dataSync()[action] = target;

        this.model.fit(tf.tensor2d([state]), qValues, { epochs: 1, verbose: 0 });
    }

    if (this.epsilon > this.epsilonMin) {
        this.epsilon *= this.epsilonDecay;
    }
}
}


function trainAgent(env, agent, episodes, batchSize) {
    for (let episode = 0; episode < episodes; episode++) {
        let state = env.reset();
        let totalReward = 0;
        const initialPortfolioValue = env.getCurrentPortfolioValue();

        while (!env.done) {
            const action = agent.selectAction(state);
            const { state: nextState, reward, done } = env.step(action);
            
            if (nextState === null) {
                break;
            }

            agent.replayBuffer.add([state, action, reward, nextState, done]);
            state = nextState;
            totalReward += reward;

            if (agent.replayBuffer.buffer.length >= batchSize) {
                const batch = agent.replayBuffer.sample(batchSize);
                agent.train(batch);
            }
        }

        const finalPortfolioValue = env.getCurrentPortfolioValue();
        const pnl = finalPortfolioValue - initialPortfolioValue;

        console.log(`Episode ${episode + 1}: Steps = ${env.currentStep}, Total Reward = ${totalReward.toFixed(2)}, PnL = $${pnl.toFixed(2)}, Final Portfolio Value = $${finalPortfolioValue.toFixed(2)}`);
    }
}
  
  // Example usage
  console.log("Initializing...");
  const data = candlestickData;
  const env = new TradingEnvironment(data);
  const agent = new DQNAgent(5, 3);  // 5 state features (OHLCV), 3 actions (Buy, Sell, Hold)
  
  console.log("Starting training process...");
  trainAgent(env, agent, 1000, 32);  // Train for 1000 episodes with batch size of 32
  console.log("Process completed.");
  



For some reason, my cahs for buying and selling is always 10000 printed to console.