text: /*
   * Copyright (c) 2020 Huawei Technologies Co.,Ltd.
   *
   * openGauss is licensed under Mulan PSL v2.
   * You can use this software according to the terms and conditions of the Mulan PSL v2.
   * You may obtain a copy of Mulan PSL v2 at:
   *
   *          http://license.coscl.org.cn/MulanPSL2
   *
   * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
   * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
   * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
   * See the Mulan PSL v2 for more details.
   */
   
   
  
  
  import math
  import random
  import torchfold
  import torch.nn.functional as F
  from torch import nn
  import torchvision.transforms as T
  import torch
  from collections import namedtuple
  from sqlSample import JoinTree
  import torch.optim as optim
  import numpy as np
  from math import log
  from itertools import count
  from PGUtils import pgrunner
  
  from ImportantConfig import Config
  # def Floss(value,targetvalue):
  #     with torch.no_grad():
  #         disl1 = (torch.abs(value-targetvalue)<1).float()
  #     return torch.mean((1-disl1)*torch.abs(value-targetvalue)*(targetvalue+1)+disl1*(value-targetvalue)*(value-targetvalue))
  def Floss(value,targetvalue):
      with torch.no_grad():
          disl1 = (torch.abs(value-targetvalue)<0.15).float()
      with torch.no_grad():
          disl2 = 1-((value>targetvalue).float()*(targetvalue>config.maxR-0.1).float())
      
      return torch.mean(disl2*((1-disl1)*torch.abs(value-targetvalue)*(targetvalue+1)+disl1*(value-targetvalue)*(value-targetvalue)))
  config = Config()
  class ENV(object):
      def __init__(self,sql,db_info,pgrunner,device,run_mode = False):
          self.sel = JoinTree(sql,db_info )
          self.sql = sql
          self.hashs = ""
          self.table_set = set([])
          self.res_table = []
          self.init_table = None
          self.planSpace = 0#0:leftDeep,1:bushy
          self.run_mode = run_mode
  
      def actionValue(self,left,right,model):
          self.sel.joinTables(left,right,fake = True)
          res_Value = self.selectValue(model)
          self.sel.total -= 1
          self.sel.aliasnames_root_set.remove(self.sel.total)
          self.sel.aliasnames_fa.pop(self.sel.left_son[self.sel.total])
          self.sel.aliasnames_fa.pop(self.sel.right_son[self.sel.total])
          return res_Value
  
      def selectValue(self,model):
          tree_state = []
          # print("--------")
          for idx in self.sel.aliasnames_root_set:
              # print(idx,idx in self.sel.aliasnames_fa,isinstance(idx,int))
              # if (not idx in self.sel.aliasnames_fa )and isinstance(idx,int):
              if (not idx in self.sel.aliasnames_fa ):
                  tree_state.append(self.sel.encode_tree_regular(model,idx))
          # if len(tree_state)!=1:
          #     print(-1)
          res = torch.cat(tree_state,dim = 0)
          # if len(self.sel.aliasnames_root_set)==len(self.join_list)+1:
          #     print(res)
          return model.logits(res,self.sel.join_matrix,len(self.sel.aliasnames_root_set)==len(self.sel.join_list)+1)
  
      def selectValueFold(self,fold):
          tree_state = []
          for idx in self.sel.aliasnames_root_set:
              # if (not idx in self.sel.aliasnames_fa )and isinstance(idx,int):
              if (not idx in self.sel.aliasnames_fa):
                  tree_state.append(self.sel.encode_tree_fold(fold,idx))
              #         res = torch.cat(tree_state,dim = 0)
          return tree_state
          return fold.add('logits',tree_state,self.sel.join_matrix)
  
  
  
      def takeAction(self,left,right):
          self.sel.joinTables(left,right)
          self.hashs += left
          self.hashs += right
          self.hashs += " "
  
      def hashcode(self):
          return self.sql.sql+self.hashs
      def allAction(self,model):
          action_value_list = []
          for one_join in self.sel.join_candidate:
              l_fa = self.sel.findFather(one_join[0])
              r_fa  =self.sel.findFather(one_join[1])
              if self.planSpace ==0:
                  flag1 = one_join[1] ==r_fa and l_fa !=one_join[0]
                  if l_fa!=r_fa and (self.sel.total == 0 or flag1):
                      action_value_list.append((self.actionValue(one_join[0],one_join[1],model),one_join))
              elif self.planSpace==1:
                  if l_fa!=r_fa:
                      action_value_list.append((self.actionValue(one_join[0],one_join[1],model),one_join))
          return action_value_list
      def reward_new(self,):
          from math import e
          if self.sel.total+1 == len(self.sel.from_table_list):
              if self.run_mode == False:
                  return self.sel.plan2Cost()/self.sel.sqlt.getDPlantecy(), True
              else:
                  return self.sel.getResult(),True
          else:
              return 0,False
  
  
  
  Transition = namedtuple('Transition',
                          ('env', 'next_value', 'this_value'))
  
  class ReplayMemory(object):
  
      def __init__(self, capacity):
          self.capacity = capacity
          self.memory = []
          self.position = 0
          self.bestJoinTreeValue = {}
      def push(self, *args):
          """Saves a transition."""
          if len(self.memory) < self.capacity:
              self.memory.append(None)
          data =  Transition(*args)
          hashv = data.env.hashcode()
          next_value = data.next_value
          this_value = data.this_value
          if hashv in self.bestJoinTreeValue and self.bestJoinTreeValue[hashv]<data.this_value:
              if self.bestJoinTreeValue[hashv]<next_value:
                  next_value = self.bestJoinTreeValue[hashv]
              this_value = self.bestJoinTreeValue[hashv]
              # import random
              # if (data.this_value < self.bestJoinTreeValue[hashv]+0.3 and random.random()>0.3):
              #     return
          else:
              self.bestJoinTreeValue[hashv]  = data.this_value
          # from math import e
          # if next_value<22220:
          # print(self.bestJoinTreeValue[hashv])
          data = Transition(data.env,self.bestJoinTreeValue[hashv],hashv)
          # else:
          #     data = Transition(data.env,next_value,this_value)
          position = self.position
          self.memory[position] = data
          self.position = (self.position + 1) % self.capacity
  
      def sample(self, batch_size):
          if len(self.memory)>batch_size:
              return random.sample(self.memory, batch_size)
          else:
              return self.memory
  
      def __len__(self):
          return len(self.memory)
      def resetMemory(self,):
          self.memory =[]
          self.position = 0
      def resetbest(self):
          self.bestJoinTreeValue = {}
  
  class DQN:
      def __init__(self,policy_net,target_net,db_info,pgrunner,device):
          self.Memory = ReplayMemory(config.memory_size)
          self.BATCH_SIZE = 1
  
          self.optimizer = optim.Adam(policy_net.parameters(),lr = 3e-4   ,betas=(0.9,0.999))
          # self.optimizer = optim.SGD(policy_net.parameters(),lr = config.learning_rate)
  
          self.steps_done = 0
          self.max_action = 25
          # self.EPS_START = 0.8
          # self.EPS_END = 0.1
          # self.EPS_DECAY = 50
          self.EPS_START = config.EPS_START
          self.EPS_END =  config.EPS_END
          self.EPS_DECAY = config.EPS_DECAY
          self.policy_net = policy_net
          self.target_net = target_net
          self.db_info = db_info
          self.device = device
          self.steps_done = 0
      def select_action(self, env, need_random = True):
  
          sample = random.random()
          if need_random:
              eps_threshold = self.EPS_END + (self.EPS_START - self.EPS_END) * \
                                        math.exp(-1. * self.steps_done / self.EPS_DECAY)
              self.steps_done += 1
          else:
              eps_threshold = -1
          action_list = env.allAction(self.policy_net)
          action_batch = torch.cat([x[0] for x in action_list],dim = 1)
          # print('select_action',sample > eps_threshold)
          # action_list
          if sample > eps_threshold:
              return action_batch,action_list[torch.argmin(action_batch,dim = 1)[0]][1],[x[1] for x in action_list]
          else:
              return action_batch,action_list[random.randint(0,len(action_list)-1)][1],[x[1] for x in action_list]
  
  
      def validate(self,val_list, tryTimes = 1):
          rewards = []
          prt = []
          mes = 0
          val_this_time = {}
          DP_cost = 0.0
          my_cost = 0.0
          allRes = {}
          import time
          startTime = time.time()
          valInfo  = {}
          for sql in val_list:
              # print('DQN')
              pg_cost = sql.getDPlantecy()
              # print('outPG_cost')
              DP_cost += pg_cost
              env = ENV(sql,self.db_info,pgrunner,self.device)
  
              for t in count():
  
                  action_list, chosen_action,all_action = self.select_action(env,need_random=False)
  
                  left = chosen_action[0]
                  right = chosen_action[1]
                  env.takeAction(left,right)
  
                  startTime = time.time()
                  reward, done = env.reward_new()
                  thisTime = time.time() -startTime
                  lalias = len(env.sel.from_table_list)
                  if done:
                      val_this_time[sql.filename] = reward
                      allRes[sql.filename] = (pg_cost,reward*pg_cost)
                      rewards.append(reward)
                      mes = mes + log(reward)
                      my_cost += reward*pg_cost
                      valInfo[sql.filename] = (lalias,thisTime,reward*pg_cost,pg_cost)
                      break
          import json
          lr = len(rewards)
          from math import e
          print("MRC",sum(rewards)/lr,"GMRL",e**(mes/lr),"SMRC",my_cost/DP_cost)
          return sum(rewards)/lr,e**(mes/lr),my_cost/DP_cost
      def validate_ind(self,val_list, tryTimes = 1):
          rewards = []
          prt = []
          mes = 0
          for sql in val_list:
              pg_cost = sql.getDPlantecy()
              env = ENV(sql,self.db_info,pgrunner,self.device)
              if (len(env.sel.from_table_list)<3):# or not env.sel.baseline.left_deep or (sql.bestLatency is not None and sql.bestLatency>config.baselineValue)) and not config.testGen:
                      rewards.append(-1.0)
                      continue
              if (len(env.sel.from_table_list)<3):
                  rewards.append(-1.0)
                  continue
              if (sql.useCost):
                  rewards.append(-1.0)
                  print('useCost')
                  continue
              for t in count():
                  action_list, chosen_action,all_action = self.select_action(env,need_random=False)
                  left = chosen_action[0]
                  right = chosen_action[1]
                  env.takeAction(left,right)
  
                  reward, done = env.reward_new()
                  if done:
                      rewards.append(reward)
                      break
          return rewards
      def optimize_model(self,):
          import time
          startTime = time.time()
          samples = self.Memory.sample(128)
          value_now_list = []
          next_value_list = []
          # if (len(samples)<64):
          #     return 0
          usecuda = True if config.usegpu == 1 else False
          # print(torchfold.path)
          fold = torchfold.Fold(cuda=usecuda)
          nowL = []
          for one_sample in samples:
              nowList = one_sample.env.selectValueFold(fold)
              nowL.append(len(nowList))
              value_now_list+=nowList
          res = fold.apply(self.policy_net, [value_now_list])[0]
          total = 0
          value_now_list = []
          next_value_list = []
          for idx,one_sample in enumerate(samples):
              value_now_list.append(self.policy_net.logits(res[total:total+nowL[idx]] , one_sample.env.sel.join_matrix ))
              next_value_list.append(min(one_sample.next_value,self.Memory.bestJoinTreeValue[one_sample.this_value]))
              total += nowL[idx]
          value_now = torch.cat(value_now_list,dim = 0)
          next_value = torch.cat(next_value_list,dim = 0)
          endTime = time.time()
          if True:
              loss = Floss(value_now,next_value)
              # loss = F.smooth_l1_loss(value_now,next_value,size_average=True)
              self.optimizer.zero_grad()
              loss.backward()
              for group in self.optimizer.param_groups:
                  for param in group["params"]:
                      if param.grad is not None:
                          param.grad.data.clamp_(-1, 1)
              self.optimizer.step()
              return loss.item()
          return None
  
          