# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.1.0
#     language: julia
#     name: julia-1.1
# ---

f_item = open("track1/item.txt", "r")
n = countlines(f_item)

f_rec_log_train = open("track1/pre_rec_data.txt", "r")
m = countlines(f_rec_log_train)

f_user_profile = open("track1/user_profile.txt", "r")
cnt_user = countlines(f_user_profile)

# +
seekstart(f_rec_log_train)
Ω = Array{Any, 1}()
for i = 1:m
    x = split(readline(f_rec_log_train), '\t')
    x = map(t->parse(Int32, t), x[1:4])
    push!(Ω, x)
end

cnt_pos, cnt_neg = 0, 0
for x in Ω
    (x[3] == -1) && (cnt_neg += 1)
    (x[3] == 1) && (cnt_pos += 1)
end
println("$cnt_pos $cnt_neg")
# -

for x in Ω[1:20]
    println(x)
end

struct User
    id::Int32
    age::Int16
    gender::Int8
    cnt_tweet::Int32
end

seekstart(f_user_profile)
U = Array{User, 1}()
for i = 1:cnt_user
    x = split(readline(f_user_profile), '\t')
    try
        x = map(t->parse(Int32, t), x[1:4])
        x[2] = (x[2] < 1950) ? 0 :
               (1950 ≤ x[2] < 2004) ? ceil((x[2] - 1950) / 2) + 1 : 
               (x[2] ≥ 2004) ? 28 : 29
    catch ArgumentError
        x[2] = "29"
        x = map(t->parse(Int32, t), x[1:4])
    end
    push!(U, User(x[1], x[2], x[3], x[4]))
end

user_dict = Dict{Int32, Int32}()
rev_user_dict = Dict{Int32, Int32}()
id = 0
for u in U
    id += 1
    user_dict[u.id] = id
    rev_user_dict[id] = u.id
    # (id <= 20) && println(rev_user_dict[id])
end

seekstart(f_item)
item_dict = Dict{Int32, Int32}()
rev_item_dict = Dict{Int32, Int32}()
for i = 1:n
    x = split(readline(f_item), '\t')
    item_dict[parse(Int32, x[1])] = i
    rev_item_dict[i] = parse(Int32, x[1])
    # (i <= 20) && println(rev_item_dict[i])
end

for x in Ω
    x[1], x[4] = user_dict[x[1]], item_dict[x[4]]
end

# +
function evaluate(p, q, bᵤ, bᵢ)
    s = 0
    for x in Ω
        r = (x[3] == 1) ? 1 : 0
        r -= sigmoid(μ + bᵤ[x[1]] + bᵢ[x[4]] + q[:,x[4]]' * p[:,x[1]])
        s += r * r
    end
    s /= m
    s += λ₁ * sum(p .* p) + λ₂ * sum(q .* q) + λ₃ * sum(bᵤ .* bᵤ) + λ₄ * sum(bᵢ .* bᵢ)
    return s
end

function sigmoid(x)
    return 1 / (1 + exp(-x))
end

λ₁, λ₂, λ₃, λ₄ = 0.0, 0.0, 0.0, 0.0
μ = cnt_pos / m
d = 40
α = 0.001
iter, max_iter = 0, 100
#p = rand(d, cnt_user)
#q = rand(d, n)
#bᵤ = rand(1, cnt_user)
#bᵢ = rand(1, n)

lst_error = evaluate(p, q, bᵤ, bᵢ)
println("$iter $lst_error")
while true
    try
        iter += 1
        for x in Ω
            r = (x[3] == 1) ? 1 : 0
            r -= sigmoid(μ + bᵤ[x[1]] + bᵢ[x[4]] + q[:,x[4]]' * p[:,x[1]])
            bᵤ[x[1]] += α * (r - λ₁ * bᵤ[x[1]])
            bᵢ[x[4]] += α * (r - λ₂ * bᵢ[x[4]])
            q[:,x[4]] += α * (r * p[:,x[1]] - λ₃ * q[:,x[4]])
            p[:,x[1]] += α * (r * q[:,x[4]] - λ₄ * p[:,x[1]])
        end
        error = evaluate(p, q, bᵤ, bᵢ)
        (iter % 1 == 0) && println("$iter $error")
        if lst_error - error < 1e-5 #|| iter >= max_iter
            break
        end
        lst_error = error
    catch InterruptException
        break
    end
end
println(lst_error)

# +
function read_from_log_test(f)
    x = split(readline(f), '\t')
    x = map(t->parse(Int32, t), x[1:4])
    return x
end

function test(f_rec_log_test, fout, key)
    m_test = countlines(f_rec_log_test)
    seekstart(f_rec_log_test)
    user_cnt, i = 0, 1
    cur = read_from_log_test(f_rec_log_test)
    while i <= m_test
        j = i + 1
        user_cnt += 1
        _Ω = Array{Any, 1}()
        push!(_Ω, cur)
        while j <= m_test
            cur = read_from_log_test(f_rec_log_test)
            (cur[1] != _Ω[1][1]) && break
            j += 1
            push!(_Ω, cur)
        end
        #(i ÷ 5000000 != j ÷ 5000000) && println(i, ' ', j)
        dict = Array{Any, 1}()
        for x in _Ω[1:j-i]
            x₁, x₂ = user_dict[x[1]], item_dict[x[2]]
            score = sigmoid(μ + bᵤ[x₁] + bᵢ[x₂] + q[:, x₂]' * p[:, x₁])
            push!(dict, (score, x[2]))
        end
        sort!(dict)
        write(fout, "$(_Ω[1][1]),")
        for k = 1:min(size(dict, 1), 3)
            t = size(dict, 1) - k + 1
            # (dict[t][1] < 0.5) && break
            write(fout, ((k == 1) ? "" : " ") * "$(dict[t][2])")
        end
        write(fout, "\n")
        i = j
    end
end

fout = open("test.csv", "w")
write(fout, "id,clicks\n")

# test(open("track1/rec_log_test_sorted.txt", "r"), fout, "Public")
test(open("track1/rec_log_test_pub.txt", "r"), fout, "Public")
test(open("track1/rec_log_test_pri.txt", "r"), fout, "Private")

close(fout)
# -



StarcoderdataJupyterScriptsDedupFiltered