import random
import numpy as np

class Asset():
  def __init__(self, cost, prob):
    """Create a new state space with given dimensions."""
    self.cost = cost
    self.prob = prob
    self.volume = []
    print("cost: ", cost)
    print("prob: ", prob)
    print("Assume volume is bound by 0 to 100")
    print("Objective: max SUM [ vol_s * prob_s - cost_s ]")

  def hill_climb(self, maximum=None, log=False):
    """Performs hill-climbing to find a solution."""
    count = 0

    self.volume=np.random.randint(0,100,len(self.prob))

    if log:
      print("Initial state: profit", self.get_profit(self.volume))

    # Continue until we reach maximum number of iterations
    while maximum is None or count < maximum:
      count += 1
      best_neighbors = []
      best_neighbor_profit = None

      for i,vol_s in enumerate(self.volume):

        for replacement in self.get_neighbors(vol_s):
          neighbor = self.volume.copy()
          neighbor[i] = replacement

          # Check if neighbor is best so far
          profit = self.get_profit(neighbor)
          if best_neighbor_profit is None or profit > best_neighbor_profit:
            best_neighbor_profit = profit
            best_neighbors = [neighbor]
          elif best_neighbor_profit == profit:
            best_neighbors.append(neighbor)

      # None of the neighbors are better than the current state
      if best_neighbor_profit <= self.get_profit(self.volume):
        return self.volume

      # Move to a highest-valued neighbor
      else:
        if log:
          print(f"Found better neighbor: cost {best_neighbor_profit}")
        self.volume = random.choice(best_neighbors)


  def random_restart(self, maximum, image_prefix=None, log=False):
    """Repeats hill-climbing multiple times."""
    best_volume = None
    best_profit = None

    # Repeat hill-climbing a fixed number of times
    for i in range(maximum):
      volume = self.hill_climb()
      profit = self.get_profit(volume)
      if best_profit is None or profit > best_profit:
        best_profit = profit
        best_volume = volume
        if log:
          print(f"{i}: Found new best state: profit {profit} with volume {volume}")
      else:
        if log:
          print(f"{i}: Found state: profit {profit}")

    return best_volume

  def get_profit(self, volume):
    profit = 0
    for i,vol_s in enumerate(self.volume):
      profit += vol_s * self.prob[i] - self.cost[i]
    return profit

  def get_neighbors(self, vol_s):
    candidates = [
      vol_s-5, vol_s+5
    ]
    neighbors = []
    for c in candidates:
      if 0 <= c < 100:
        neighbors.append(c)
    return neighbors


if __name__ == "__main__":

  num_assets=10
  s = Asset(np.random.randint(0, 10, num_assets),
            np.random.rand(num_assets))

  s.random_restart(10, log=True)