from matplotlib import pyplot as plt
import numpy as np
from collections import OrderedDict
import os
from tqdm import tqdm
import sys
from typing import Dict
import typing
import torch
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--base_model", type=str, required=True)
parser.add_argument("--lora_load", type=str, required=True)
parser.add_argument("--lora_schedule", type=str, default="const", choices=["const", "linear", "poly"])
parser.add_argument("--lora_alpha", type=float, default=None)
parser.add_argument("--lora_beta", type=float, default=None)
parser.add_argument("--lora_strategy", type=str, default=None)
parser.add_argument("--plot", action="store_true")
parser.add_argument("--save", required=True, type=str)
args = parser.parse_args()

def f_const(alpha):
    def inner(x):
        return alpha
    return inner
def f_linear(alpha, beta):
    def inner(x):
        return alpha+x*(beta-alpha)
    return inner

def f_polyline(*args):
    args = sorted(args)
    xs = []
    ys = []
    xys = args
    for x, y in args:
        xs.append(x)
        ys.append(y)
    def inner(x):
        if(x<=xs[0]):
            return ys[0]
        if(x>=xs[-1]):
            return ys[-1]
        for i, x1 in enumerate(xs):
            x0 = xs[i-1]
            if(x0<=x and x<=x1):
                y0 = ys[i-1]
                y1 = ys[i]
                n = (x-x0)/(x1-x0)
                y = y0*(1-n)+y1*n
                return y
        return ys[-1]
    return inner


with torch.no_grad():
    base = args.base_model
    lora = args.lora_load
    w: Dict[str, torch.Tensor] = torch.load(base, map_location="cpu")
    w_lora: Dict[str, torch.Tensor] = torch.load(lora, map_location="cpu")
    out: typing.OrderedDict[str, torch.Tensor] = OrderedDict()

    tuned_weights = []

    lora_rank = None
    listkeys = list(w.keys())
    for key in listkeys:
        if key.endswith(".weight"):
            pref = key[:-len(".weight")]
            key_A = pref+".lora_A"
            key_B = pref+".lora_B"
            if(key_A in w_lora and key_B in w_lora):
                print("Weight %s will be merged with lora"%pref)
                tuned_weights.append(key)

                if(lora_rank is None):
                    lora_A = w_lora[key_A]
                    lora_B = w_lora[key_B]
                    lora_rank = lora_A.shape[0]
                    assert lora_rank == lora_B.shape[1]
    
    
    if(args.lora_schedule == "const"):
        alpha = args.lora_alpha
        if(alpha is None):
            print("assuming lora_alpha = lora_rank = %s"%lora_rank)
            alpha = lora_rank
        f_schedule = f_const(alpha)
    elif(args.lora_schedule == "poly"):
        xys = []
        for i in args.lora_strategy.split(","):
            xy = [float(_) for _ in i.split(":")]
            xys.append(xy)
        f_schedule = f_polyline(*xys)
    else:
        alpha = args.lora_alpha
        beta = args.lora_beta
        if(alpha is None):
            print("assuming lora_alpha = lora_rank = %s"%lora_rank)
            alpha = lora_rank
        if(beta is None):
            print("assuming beta = lora_rank = %s"%lora_rank)
            beta = lora_rank
        f_schedule = f_linear(alpha, beta)
    
    weight2alpha = {}
    for idx, i in enumerate(tuned_weights):
        # assuming weights are fetched in depth order
        depth = idx
        if(len(tuned_weights)>1):
            depth = depth/(len(tuned_weights)-1)
        else:
            depth = 0.5
        alpha = f_schedule(depth)
        print("%s: alpha=%.3f, alpha/rank=%.2f"%(i, alpha, alpha/lora_rank))
        weight2alpha[i] = alpha
    if(args.plot):
        tot = 32
        xs = np.linspace(0, tot-1, tot)
        ys = [f_schedule(x/(tot-1)) for x in xs]
        plt.plot(xs, ys)
        plt.savefig(args.save+".png")
        print("plot", args.save+".png")
    for key in tqdm(listkeys):
        if(key in tuned_weights):
            pref = key[:-len(".weight")]
            key_A = pref+".lora_A"
            key_B = pref+".lora_B"
            
            lora_A = w_lora[key_A]
            lora_B = w_lora[key_B]
            lora_rank = lora_A.shape[0]
            assert lora_rank == lora_B.shape[1]

            delta_w = (lora_B @ lora_A)*weight2alpha[key]/lora_rank
            new_w = w[key]+delta_w
            out_w = new_w.to(device=w[key].device, dtype=w[key].dtype, copy=True)
            out[key] = out_w
            del w[key]
        else:
            # print("%s = original"%(key))
            out[key] = w[key].clone()
            del w[key]
    # assert list(out.keys()) == list(w.keys())
    # for key in out.keys():
    #     assert out[key].shape == w[key].shape
    #     assert out[key].dtype == w[key].dtype
    
    
    torch.save(out, args.save)
    print("saved", args.save)