import gradio as gr
import json
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from collections import namedtuple
import datetime
import calendar
from dateutil.relativedelta import relativedelta

def time2date(input):
    chu_ky_thoi_gian = input['CHU KỲ THỜI GIAN']
    thu = input['THỨ']
    ngay = input['NGÀY']
    tuan = input['TUẦN']
    thang = input['THÁNG']
    quy = input['QUÝ']
    nam = input['NĂM']
    current_date = datetime.date.today()
    output = 'error'
    if chu_ky_thoi_gian=='ngày':
# hôm kia
        if ngay=='hôm kia':
            output = current_date - datetime.timedelta(days=2)
# hôm qua
        elif ngay=='hôm qua':
            output = current_date - datetime.timedelta(days=1)
# hôm nay
        elif ngay=='hôm nay':
            output = current_date
# ngày mai
        elif ngay=='mai':
            output = current_date + datetime.timedelta(days=1)
# ngày kia
        elif ngay=='kia':
            output = current_date + datetime.timedelta(days=2)
# đầu
        elif ngay=='đầu':
            if thang=='trước':
                needed_thang = 12 if current_date.month==1 else current_date.month-1
                needed_nam = current_date.year-1 if current_date.month==1 else current_date.year
                output = datetime.date(
                    needed_nam,
                    needed_thang,
                    1
                )
            elif thang=='này':
                needed_thang = current_date.month
                needed_nam = current_date.year
                output = datetime.date(
                    needed_nam,
                    needed_thang,
                    1
                )
            elif thang=='sau':
                needed_thang = 1 if current_date.month==12 else current_date.month+1
                needed_nam = current_date.year+1 if current_date.month==12 else current_date.year
                output = datetime.date(
                    needed_nam,
                    needed_thang,
                    1
                )
            else:
                print('ngày đầu')
                output = current_date
# cuối
        elif ngay=='cuối':
            if thang=='trước':
                needed_thang = 12 if current_date.month==1 else current_date.month-1
                needed_nam = current_date.year-1 if current_date.month==1 else current_date.year
                output = datetime.date(
                    needed_nam,
                    needed_thang,
                    calendar.monthrange(needed_nam, needed_thang)[1]
                )
            elif thang=='này':
                needed_thang = current_date.month
                needed_nam = current_date.year
                output = datetime.date(
                    needed_nam,
                    needed_thang,
                    calendar.monthrange(needed_nam, needed_thang)[1]
                )
            elif thang=='sau':
                needed_thang = 1 if current_date.month==12 else current_date.month+1
                needed_nam = current_date.year+1 if current_date.month==12 else current_date.year
                output = datetime.date(
                    needed_nam,
                    needed_thang,
                    calendar.monthrange(needed_nam, needed_thang)[1]
                )
            else:
                print('ngày cuối')
                output = current_date
# xxx
        elif ngay.isdigit():
            ngay = int(ngay)
            if thang.isdigit():
                thang = int(thang)
                if nam.isdigit():
                    nam = int(nam)
                    try:
                        output = datetime.date(nam, thang, ngay)
                    except ValueError:
                        print("ngày xxx tháng yyy nam xxx")
                        output = current_date
                else:
                    try:
                        output = datetime.date(current_date.year, thang, ngay)
                    except ValueError:
                        print("ngày xxx tháng yyy nam xxx")
                        output = current_date
            else:
                try:
                    output = datetime.date(current_date.year, current_date.month, ngay)
                except ValueError:
                    print('ngày xxx')
                    output = current_date
        elif ngay=='khác':
# thứ hai tuần sau
# thứ hai tuần này
# thứ hai tuần trước
# thứ hai tuần gần nhất
# thứ hai
            if thu != "khác":
                if tuan=='khác' or tuan=='gần nhất':
                    current_thu = current_date.weekday()+2
                    input_thu = 8 if thu=='chủ nhật' else int(thu)
                    if input_thu < current_thu:
                        daydelta = current_thu - input_thu
                        output = current_date - datetime.timedelta(days=daydelta)
                    elif input_thu == current_thu:
                        output = current_date - datetime.timedelta(days=7)
                    else:
                        daydelta = input_thu - current_thu
                        output = current_date - datetime.timedelta(days=7-daydelta)
                elif tuan=='này':
                    current_thu = current_date.weekday()+2
                    input_thu = 8 if thu=='chủ nhật' else int(thu)
                    if input_thu < current_thu:
                        daydelta = current_thu - input_thu
                        output = current_date - datetime.timedelta(days=daydelta)
                    elif input_thu == current_thu:
                        output = current_date
                    else:
                        daydelta = input_thu - current_thu
                        output = current_date + datetime.timedelta(days=daydelta)
                elif tuan=='trước':
                    current_thu = current_date.weekday()+2
                    input_thu = 8 if thu=='chủ nhật' else int(thu)
                    if input_thu < current_thu:
                        daydelta = current_thu - input_thu + 7
                        output = current_date - datetime.timedelta(days=daydelta)
                    elif input_thu == current_thu:
                        output = current_date - datetime.timedelta(days=7)
                    else:
                        daydelta = input_thu - current_thu
                        output = current_date - datetime.timedelta(days=7-daydelta)
                elif tuan=='sau':
                    current_thu = current_date.weekday()+2
                    input_thu = 8 if thu=='chủ nhật' else int(thu)
                    if input_thu < current_thu:
                        daydelta = 7 - (current_thu - input_thu)
                        output = current_date + datetime.timedelta(days=daydelta)
                    elif input_thu == current_thu:
                        output = current_date + datetime.timedelta(days=7)
                    else:
                        daydelta = input_thu - current_thu
                        output = current_date + datetime.timedelta(days=7+daydelta)
            elif thu=='khác':
                current_thu = current_date.weekday()
                if tuan=='trước' or tuan=='gần nhất': # ngày cuối tuần
                    daydelta = current_thu+1
                    output = current_date - datetime.timedelta(days=daydelta)
                elif tuan=='này':
                    daydelta = 6 - current_thu
                    if daydelta > 0:
                        output = current_date + datetime.timedelta(days=daydelta)
                    else:
                        output = current_date
                elif tuan=='sau':
                    daydelta = 13 - current_thu
                    output = current_date + datetime.timedelta(days=daydelta)
                else:
                    print('ngày khác thứ khác tuần khác')
                    output = current_date
        else: # return default for ngay
            output = current_date
            print('Error chu ky thoi gian: ngay')

    elif chu_ky_thoi_gian=='tháng': # ngày cuối tháng
        if thu != 'khác' or ngay != 'khác' or tuan != 'khác':
            print('Invalid')
            needed_thang = 12 if current_date.month==1 else current_date.month-1
            needed_nam = current_date.year-1 if current_date.month==1 else current_date.year
            output = datetime.date(
                needed_nam,
                needed_thang,
                calendar.monthrange(needed_nam, needed_thang)[1]
            )

        elif thang=='trước':
            needed_thang = 12 if current_date.month==1 else current_date.month-1
            needed_nam = current_date.year-1 if current_date.month==1 else current_date.year
            output = datetime.date(
                needed_nam,
                needed_thang,
                calendar.monthrange(needed_nam, needed_thang)[1]
            )
        elif thang=='này':
            needed_thang = current_date.month
            needed_nam = current_date.year
            output = datetime.date(
                needed_nam,
                needed_thang,
                calendar.monthrange(needed_nam, needed_thang)[1]
            )
        elif thang=='sau':
            needed_thang = 1 if current_date.month==12 else current_date.month+1
            needed_nam = current_date.year+1 if current_date.month==12 else current_date.year
            output = datetime.date(
                needed_nam,
                needed_thang,
                calendar.monthrange(needed_nam, needed_thang)[1]
            )
        elif thang.isdigit():
            thang = int(thang)
            if nam.isdigit():
                nam = int(nam)
                output = datetime.date(
                    int(nam),
                    int(thang),
                    calendar.monthrange(int(nam), int(thang))[1]
                )
            else:
                if thang > current_date.month:
                    output = datetime.date(
                        current_date.year-1,
                        thang,
                        calendar.monthrange(current_date.year-1, thang)[1]
                    )
                else:
                    output = datetime.date(
                        current_date.year,
                        thang,
                        calendar.monthrange(current_date.year, thang)[1]
                    )
        elif thang=='khác':
            needed_thang = 12 if current_date.month==1 else current_date.month-1
            needed_nam = current_date.year-1 if current_date.month==1 else current_date.year
            output = datetime.date(
                needed_nam,
                needed_thang,
                calendar.monthrange(needed_nam, needed_thang)[1]
            )
        else: # ngay cuoi thang truoc
            needed_thang = 12 if current_date.month==1 else current_date.month-1
            needed_nam = current_date.year-1 if current_date.month==1 else current_date.year
            output = datetime.date(
                needed_nam,
                needed_thang,
                calendar.monthrange(needed_nam, needed_thang)[1]
            )
            print('Error chu ky thoi gian: thang')
    elif chu_ky_thoi_gian=='quý':
        if quy in ['1', '2', '3', '4', 'I', 'II', 'III', 'IV']:
            output = "Valid"
            output = current_date
            # print('Chu ky thoi gian: quy ' + quy)
        else:
            output = "Invalid"
            output = current_date
            # print('Error chu ky thoi gian: quy')
    elif chu_ky_thoi_gian=='năm':
        if nam.isdigit():
            output = "Valid"
            output = current_date
            # print("Chu ky thoi gian: nam " + nam)
        else:
            output = "Invalid"
            output = current_date
            # print('Error chu ky thoi gian: năm')
    elif chu_ky_thoi_gian=='khác':
        output = "Valid"
        output = current_date
    else:
        output = current_date
    return output

fields = ['device', 'model_name', 'max_source_length', 'max_target_length', 'beam_size']
params = namedtuple('params', field_names=fields)

args = params(
    device="cuda" if torch.cuda.is_available() else "cpu",
    model_name='facebook/mbart-large-50-many-to-many-mmt',
    max_source_length=256,
    max_target_length=256,
    beam_size=1
)

model = AutoModelForSeq2SeqLM.from_pretrained(
    "Huy1432884/db_retrieval",
    use_auth_token="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"
)
model.to(args.device)
model.eval()

if "mbart" in args.model_name.lower():
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_name, src_lang="vi_VN", tgt_lang="vi_VN"
    )
else:
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)

with open("output2url.json") as f:
    output2url = json.loads(list(f)[0])


def text_analysis(text):
    text = text.lower()

    inputs = tokenizer(
        [text],
        text_target=None,
        padding="longest",
        max_length=args.max_source_length,
        truncation=True,
        return_tensors="pt",
    )

    for k, v in inputs.items():
        inputs[k] = v.to(args.device)


    if "mbart" in args.model_name:
        inputs["forced_bos_token_id"] = tokenizer.lang_code_to_id["vi_VN"]

    outputs = model.generate(
        **inputs,
        max_length=args.max_target_length,
        num_beams=args.beam_size,
        early_stopping=True,
    )

    output_sentences = tokenizer.batch_decode(
        outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True
    )
    out = json.loads("{" + output_sentences[0] + "}")
    if out['LOẠI BIỂU ĐỒ']=='dashboard':
        if out['CHU KỲ THỜI GIAN']!='tháng':
            chu_ky_in = 'ngày'
        else:
            chu_ky_in = 'tháng'
        out['CHU KỲ THỜI GIAN']='ngày' if out['CHU KỲ THỜI GIAN'] not in ['ngày', 'tháng'] else out['CHU KỲ THỜI GIAN']
        check_dashboard = out['ĐƠN VỊ']+"_"+chu_ky_in
        out['DB URL'] = output2url[check_dashboard]
        out['DATE'] = time2date(out)
        out['FINAL URL'] = "https://vsds.viettel.vn"+ out['DB URL'] + "?toDate=" + str(out['DATE']).replace("-", "").replace("-", "")
        show = {i: out[i] for i in ['LOẠI BIỂU ĐỒ', 'ĐƠN VỊ', 'CHU KỲ THỜI GIAN', 'DB URL', 'DATE', 'FINAL URL']}
    elif out['LOẠI BIỂU ĐỒ']=='biểu đồ':
        show = {i: out[i] for i in ['LOẠI BIỂU ĐỒ', 'ĐƠN VỊ', 'CHU KỲ THỜI GIAN']}
    else:
        show = out
    return show

demo = gr.Interface(
    text_analysis,
    gr.Textbox(placeholder="Enter sentence here..."),
    ["json"],
    examples=[
        ["Mở dashboard vtc ngày hôm qua"],
        ["Mở biểu đồ cột td ngày này"],
        ["Hãy mở biểu đồ cơ cấu của tập đoàn trong ngày hôm nay"],
        ["Tháng này, vtc cần tôi mở biểu đồ rank để cập nhật danh sách khách hàng"],
        ["Các thông số NAT ngày hôm qua đã được ghi nhận trên đát bọt"],
        ["Hôm nay hãy mở của Viettel tt không gian mạng Viettel vtcc để kiểm tra"],
        ["Mở DB CTM ngày gốc"],
        ["Tôi đã sử dụng Dashboard để truy cập thông tin qti vào ngày hôm nay"],
        ["Trưởng phòng đã ra lệnh mở biểu đồ kết hợp đường và cột cho toàn tập đoàn vào hôm nay"]
    ],
)


demo.launch()