Add export function and rewrite code